rohityadav1993 commented on code in PR #19120: URL: https://github.com/apache/pinot/pull/19120#discussion_r4306896836
########## pinot-core/src/main/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperator.java: ########## @@ -0,0 +1,575 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.core.operator.query; + +import com.google.common.base.CaseFormat; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.PriorityQueue; +import java.util.Set; +import java.util.function.Function; +import java.util.stream.Collectors; +import javax.annotation.Nullable; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.common.request.context.OrderByExpressionContext; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.common.utils.HashUtil; +import org.apache.pinot.core.common.BlockValSet; +import org.apache.pinot.core.common.Operator; +import org.apache.pinot.core.common.RowBasedBlockValueFetcher; +import org.apache.pinot.core.operator.BaseOperator; +import org.apache.pinot.core.operator.BaseProjectOperator; +import org.apache.pinot.core.operator.BitmapDocIdSetOperator; +import org.apache.pinot.core.operator.ColumnContext; +import org.apache.pinot.core.operator.ExecutionStatistics; +import org.apache.pinot.core.operator.ExplainAttributeBuilder; +import org.apache.pinot.core.operator.ProjectionOperator; +import org.apache.pinot.core.operator.ProjectionOperatorUtils; +import org.apache.pinot.core.operator.blocks.ValueBlock; +import org.apache.pinot.core.operator.blocks.results.SelectionResultsBlock; +import org.apache.pinot.core.operator.transform.TransformOperator; +import org.apache.pinot.core.operator.transform.function.TransformFunctionFactory; +import org.apache.pinot.core.query.request.context.QueryContext; +import org.apache.pinot.core.query.selection.SelectionOperatorUtils; +import org.apache.pinot.core.query.utils.OrderByComparatorFactory; +import org.apache.pinot.segment.spi.IndexSegment; +import org.apache.pinot.segment.spi.datasource.DataSource; +import org.apache.pinot.spi.query.QueryScanCostContext; +import org.roaringbitmap.RoaringBitmap; + + +/// Lazy, incremental selection ORDER BY operator for segments that are physically sorted on the first order-by column. +/// +/// Unlike [SelectionOrderByOperator] (which materializes the segment's whole top-K in a single block) this +/// operator emits one globally-sorted [SelectionResultsBlock] per [#getNextBlock()] call and returns +/// `null` when the segment is exhausted, so that a downstream k-way-merge combine operator can pull from many +/// segments lazily and stop early. It relies on the underlying project operator iterating the first order-by column in +/// the query order (the caller must guarantee `projectOperator.isCompatibleWith(DocIdOrder.fromAsc(asc))`). +/// +/// Every instance emits at least one block before `null`: a segment matching no rows emits a single empty block +/// carrying the [DataSchema], so consumers never have to reconstruct a schema the segment already knows. +/// +/// It runs in one of two emission modes: +/// +/// - **No tail to sort** (`numSortedExpressions == numOrderByExpressions`, e.g. `ORDER BY sorted`): +/// rows already arrive from the project operator in final order, so each call emits the next project block +/// (trimmed to the remaining `limit + offset` budget). +/// - **Tail to sort** (`numSortedExpressions < numOrderByExpressions`, e.g. +/// `ORDER BY sorted, other`): +/// each call reads forward until the first order-by value changes (a primary-value "run"), retains the run's top +/// `limit + offset` rows by the full comparator, and emits them sorted. This bounds the in-memory run buffer to +/// `limit + offset` rows even when the first order-by column is near-constant (very low cardinality). +/// +/// Like [SelectionOrderByOperator] it preserves the two-phase projection optimization: when there are output +/// expressions that are not order-by expressions, the forward scan only fetches the order-by expressions plus the +/// document id, and the non-order-by expressions are fetched in a second pass over the retained document ids of each +/// emitted block. +/// +/// This operator is stateful across [#getNextBlock()] calls and is **not** thread-safe; a single consumer +/// must drive it. +public class StreamingSelectionOrderByOperator extends BaseOperator<SelectionResultsBlock> { + private static final String EXPLAIN_NAME = "SELECT_ORDERBY_STREAMING"; + + private final IndexSegment _indexSegment; + private final QueryContext _queryContext; + private final boolean _nullHandlingEnabled; + /// Deduped order-by expressions followed by output expressions from SelectionOperatorUtils.extractExpressions() + private final List<ExpressionContext> _expressions; + private final BaseProjectOperator<?> _projectOperator; + private final List<OrderByExpressionContext> _orderByExpressions; + private final ColumnContext[] _orderByColumnContexts; + private final int _numExpressions; + private final int _numOrderByExpressions; + private final int _numRowsToKeep; + /// Whether there are output expressions that are not order-by expressions (requires the two-phase fetch) + private final boolean _twoPhase; + /// Whether the order-by has an unsorted tail that must be sorted in memory per run + private final boolean _tailToSort; + /// Expressions fetched during the forward scan: order-by expressions only when two-phase, otherwise all expressions + private final List<ExpressionContext> _phase1Expressions; + private final int _numPhase1Columns; + private final Comparator<Object[]> _comparator; + /// Compares only the first order-by column; used to detect primary-value run boundaries + private final Comparator<Object[]> _primaryComparator; + /// Pre-allocated run heap (cleared and reused each nextRun() call to avoid per-run allocation) + private final Comparator<Object[]> _reversedComparator; + private final PriorityQueue<Object[]> _runHeap; + + // Pre-computed invariants for the two-phase fetch (null when single-phase) + private final List<ExpressionContext> _nonOrderByExpressions; + private final Map<String, DataSource> _phase2DataSourceMap; + private final int _phase2NumColumns; + + /// Lazily built and cached; for two-phase it requires the transform operator's result column contexts + private DataSchema _dataSchema; + /// Source column contexts for the phase-2 expressions; built only on the zero-match schema path (null otherwise) + private Map<String, ColumnContext> _phase2ColumnContextMap; + + // Forward-scan cursor state (used by the tail-to-sort mode) + private ValueBlock _currentBlock; + private RowBasedBlockValueFetcher _currentFetcher; + private int[] _currentDocIds; + private RoaringBitmap[] _currentNullBitmaps; + private int _currentNumDocs; + private int _currentPos; + /// One-row lookahead: the first row of the next run, stashed when a run boundary is crossed + private Object[] _pendingRow; + private boolean _projectExhausted; + + private boolean _exhausted; + /// Whether any block has been emitted yet; drives the zero-match segment's one schema-carrying empty block + private boolean _emittedAnyBlock; + private int _numRowsEmitted; + private int _numDocsScanned = 0; + private long _numEntriesScannedPostFilter = 0; + + public StreamingSelectionOrderByOperator(IndexSegment indexSegment, QueryContext queryContext, + List<ExpressionContext> expressions, BaseProjectOperator<?> projectOperator, int numSortedExpressions) { + _indexSegment = indexSegment; + _queryContext = queryContext; + _nullHandlingEnabled = queryContext.isNullHandlingEnabled(); + _expressions = expressions; + _projectOperator = projectOperator; + + _orderByExpressions = queryContext.getOrderByExpressions(); + assert _orderByExpressions != null; + _numExpressions = expressions.size(); + _numOrderByExpressions = _orderByExpressions.size(); + _orderByColumnContexts = new ColumnContext[_numOrderByExpressions]; + for (int i = 0; i < _numOrderByExpressions; i++) { + ExpressionContext expression = _orderByExpressions.get(i).getExpression(); + _orderByColumnContexts[i] = _projectOperator.getResultColumnContext(expression); + } + + _numRowsToKeep = queryContext.getOffset() + queryContext.getLimit(); + _twoPhase = _numExpressions > _numOrderByExpressions; + _tailToSort = numSortedExpressions < _numOrderByExpressions; + _comparator = + OrderByComparatorFactory.getComparator(_orderByExpressions, _orderByColumnContexts, _nullHandlingEnabled); + // The first order-by column is the physically sorted column, so it never contains nulls on this path; comparing + // only index 0 is enough to detect when one primary-value run ends and the next begins. + _primaryComparator = + OrderByComparatorFactory.getComparator(_orderByExpressions, _orderByColumnContexts, _nullHandlingEnabled, 0, 1); + _reversedComparator = _comparator.reversed(); + _runHeap = new PriorityQueue<>( + Math.min(_numRowsToKeep, SelectionOperatorUtils.MAX_ROW_HOLDER_INITIAL_CAPACITY), _reversedComparator); + + if (_twoPhase) { + _phase1Expressions = new ArrayList<>(_numOrderByExpressions); + for (OrderByExpressionContext orderByExpression : _orderByExpressions) { + _phase1Expressions.add(orderByExpression.getExpression()); + } + _nonOrderByExpressions = _expressions.subList(_numOrderByExpressions, _numExpressions); + Set<String> columns = new HashSet<>(); + for (ExpressionContext expressionContext : _nonOrderByExpressions) { + expressionContext.getColumns(columns); + } + _phase2NumColumns = columns.size(); + _phase2DataSourceMap = new HashMap<>(); + for (String column : columns) { + _phase2DataSourceMap.put(column, _indexSegment.getDataSource(column, _queryContext.getSchema())); + } + } else { + _phase1Expressions = _expressions; + _nonOrderByExpressions = null; + _phase2NumColumns = 0; + _phase2DataSourceMap = null; + // Single-phase: all output expressions are order-by expressions, so their types are known up front. + _dataSchema = buildSinglePhaseDataSchema(); + } + _numPhase1Columns = _phase1Expressions.size(); + } + + @Override + protected SelectionResultsBlock getNextBlock() { + if (_exhausted) { + return null; + } + List<Object[]> rows = _tailToSort ? nextRun() : nextSortedRows(); + if (rows == null || rows.isEmpty()) { + _exhausted = true; + if (_emittedAnyBlock) { + return null; + } + // Zero-match segment: emit one empty block carrying the schema before signalling end-of-stream, so consumers + // never have to reconstruct a schema the segment already knows. Two-phase leaves _dataSchema null until the + // first fetch, which never happens here, so build it from column metadata alone. + if (_twoPhase && _dataSchema == null) { + _dataSchema = buildTwoPhaseDataSchema(this::resolveResultColumnContext); + } + assert _dataSchema != null; + return new SelectionResultsBlock(_dataSchema, List.of(), _comparator, _queryContext); + } + _emittedAnyBlock = true; + if (_twoPhase) { + fetchNonOrderByColumns(rows); + } + // Single-phase builds the schema in the constructor; two-phase builds it during fetchNonOrderByColumns above. + assert _dataSchema != null; + return new SelectionResultsBlock(_dataSchema, rows, _comparator, _queryContext); + } + + /// No-tail-to-sort mode: the project operator already returns rows in final order, so emit the next project block, + /// trimmed to the remaining `limit + offset` budget. Returns `null` when exhausted, and otherwise a + /// non-empty list -- empty project blocks are skipped and the row budget is checked up front, so "no rows" always + /// means "no more rows". + @Nullable + private List<Object[]> nextSortedRows() { + int remaining = _numRowsToKeep - _numRowsEmitted; + if (remaining <= 0) { + return null; + } + ValueBlock valueBlock = nextNonEmptyBlock(); + if (valueBlock == null) { + return null; + } + int numDocsFetched = valueBlock.getNumDocs(); + BlockValSet[] blockValSets = new BlockValSet[_numPhase1Columns]; + for (int i = 0; i < _numPhase1Columns; i++) { + blockValSets[i] = valueBlock.getBlockValueSet(_phase1Expressions.get(i)); + } + RowBasedBlockValueFetcher blockValueFetcher = new RowBasedBlockValueFetcher(blockValSets); + int[] docIds = _twoPhase ? valueBlock.getDocIds() : null; + RoaringBitmap[] nullBitmaps = null; + if (_nullHandlingEnabled) { + nullBitmaps = new RoaringBitmap[_numPhase1Columns]; + for (int i = 0; i < _numPhase1Columns; i++) { + nullBitmaps[i] = blockValSets[i].getNullBitmap(); + } + } + _numDocsScanned += numDocsFetched; + _numEntriesScannedPostFilter += (long) numDocsFetched * _projectOperator.getNumColumnsProjected(); + reportScanCost(numDocsFetched, (long) numDocsFetched * _projectOperator.getNumColumnsProjected()); + + // Rows arrive sorted; we only need the first 'remaining' of them globally. + int numRows = Math.min(numDocsFetched, remaining); + List<Object[]> rows = new ArrayList<>(numRows); + for (int i = 0; i < numRows; i++) { + rows.add(materializeRow(blockValueFetcher, docIds, nullBitmaps, i)); + } + _numRowsEmitted += rows.size(); + return rows; + } + + /// Tail-to-sort mode: read forward until the first order-by value changes, retain the run's top + /// `limit + offset` rows by the full comparator, and return them sorted. Returns `null` when + /// exhausted. + @Nullable + private List<Object[]> nextRun() { + int remaining = _numRowsToKeep - _numRowsEmitted; + if (remaining <= 0) { + return null; + } + if (_pendingRow == null) { + _pendingRow = nextRow(); + if (_pendingRow == null) { + return null; + } + } + PriorityQueue<Object[]> runHeap = _runHeap; + runHeap.clear(); + Object[] runFirstRow = _pendingRow; + SelectionOperatorUtils.addToPriorityQueue(_pendingRow, runHeap, _numRowsToKeep); + _pendingRow = null; + Object[] row; + while ((row = nextRow()) != null) { + if (_primaryComparator.compare(row, runFirstRow) == 0) { + SelectionOperatorUtils.addToPriorityQueue(row, runHeap, _numRowsToKeep); + } else { + // Run boundary: this row starts the next run, keep it for the next call. + _pendingRow = row; + break; + } + } + List<Object[]> rows = drainAscending(runHeap); + // A segment never contributes more than 'limit + offset' rows to the global result, and they are a prefix of its + // local sorted order, so cap the total emitted across runs at the remaining budget (the rows are ascending, keep + // the smallest 'remaining'). + if (rows.size() > remaining) { + rows = rows.subList(0, remaining); + } + _numRowsEmitted += rows.size(); + return rows; + } + + /// Pulls the next row of the forward scan (across project blocks), materialized as an + /// `Object[_numExpressions]`. For two-phase the document id is stashed at index + /// `_numOrderByExpressions` (overwritten in the second pass). Returns `null` when the project operator + /// is exhausted. + @Nullable + private Object[] nextRow() { + if (_currentBlock == null || _currentPos >= _currentNumDocs) { + if (_projectExhausted) { + return null; + } + _currentBlock = nextNonEmptyBlock(); + if (_currentBlock == null) { + _projectExhausted = true; + return null; + } + BlockValSet[] blockValSets = new BlockValSet[_numPhase1Columns]; + for (int i = 0; i < _numPhase1Columns; i++) { + blockValSets[i] = _currentBlock.getBlockValueSet(_phase1Expressions.get(i)); + } + _currentFetcher = new RowBasedBlockValueFetcher(blockValSets); + _currentNumDocs = _currentBlock.getNumDocs(); + _currentDocIds = _twoPhase ? _currentBlock.getDocIds() : null; + if (_nullHandlingEnabled) { + _currentNullBitmaps = new RoaringBitmap[_numPhase1Columns]; + for (int i = 0; i < _numPhase1Columns; i++) { + _currentNullBitmaps[i] = blockValSets[i].getNullBitmap(); + } + } + _currentPos = 0; + _numDocsScanned += _currentNumDocs; + _numEntriesScannedPostFilter += (long) _currentNumDocs * _projectOperator.getNumColumnsProjected(); + reportScanCost(_currentNumDocs, (long) _currentNumDocs * _projectOperator.getNumColumnsProjected()); + } + int rowId = _currentPos++; + return materializeRow(_currentFetcher, _currentDocIds, _currentNullBitmaps, rowId); + } + + /// Pulls the next project block carrying documents, skipping any that carry none. Returns `null` only when + /// the project operator is exhausted; an empty block means "nothing in this batch", not "end of segment", and + /// treating one as exhaustion would truncate the scan. + /// + /// No reachable doc-id-set operator emits an empty block today, so this normally returns on its first pull. The + /// skip stays because that is a convention of those implementations, not something + /// [org.apache.pinot.core.operator.BaseProjectOperator] enforces. + @Nullable + private ValueBlock nextNonEmptyBlock() { + while (true) { + ValueBlock valueBlock = _projectOperator.nextBlock(); + if (valueBlock == null) { + return null; + } + if (valueBlock.getNumDocs() > 0) { + return valueBlock; + } + } + } + + /// Materializes a single phase-1 row (deep-copied out of the value block buffers) from the given fetcher. + private Object[] materializeRow(RowBasedBlockValueFetcher fetcher, @Nullable int[] docIds, + @Nullable RoaringBitmap[] nullBitmaps, int rowId) { + Object[] row = new Object[_numExpressions]; + fetcher.getRow(rowId, row, 0); + if (_twoPhase) { + row[_numOrderByExpressions] = docIds[rowId]; + } + if (_nullHandlingEnabled) { + for (int colId = 0; colId < _numPhase1Columns; colId++) { + if (nullBitmaps[colId] != null && nullBitmaps[colId].contains(rowId)) { + row[colId] = null; + } + } + } + return row; + } + + /// Drains a max-heap (created with the reversed comparator) into an ascending list. The second pass fills + /// non-order-by values in place, so the list must at least be settable. + /// + /// It is a growable `ArrayList` rather than a fixed-size view for defensive reasons: `SelectionResultsBlock` is a + /// shared type, and `SelectionOperatorUtils.mergeWithoutOrdering()` structurally adds to the row list of the block + /// it merges into. No current consumer of this operator takes that path - the streaming combine only reads the + /// rows - so this is hardening against a future one, not a fix for a reachable failure. + private List<Object[]> drainAscending(PriorityQueue<Object[]> heap) { + int numRows = heap.size(); + // One exactly-sized backing array and no intermediate copy: drain in the heap's own (descending) order, then + // flip in place. + List<Object[]> rows = new ArrayList<>(numRows); + for (int i = 0; i < numRows; i++) { + rows.add(heap.poll()); + } + Collections.reverse(rows); + return rows; + } + + /// Second pass of the two-phase fetch: fills the non-order-by expression values for the rows of a single emitted + /// block. + /// The rows keep their final (comparator) order; the fill iterates a document-id-sorted view that shares the same row + /// instances, mirroring [SelectionOrderByOperator#computePartiallyOrdered()]. + private void fetchNonOrderByColumns(List<Object[]> rows) { + int numRows = rows.size(); + RoaringBitmap docIds = new RoaringBitmap(); + for (Object[] row : rows) { + docIds.add((int) row[_numOrderByExpressions]); + } + // Document-id-sorted view sharing the same row instances (the bitmap returns docIds in ascending order). + List<Object[]> rowsByDocId = new ArrayList<>(rows); + rowsByDocId.sort(Comparator.comparingInt(o -> (int) o[_numOrderByExpressions])); + + BitmapDocIdSetOperator docIdOperator = BitmapDocIdSetOperator.ascending(docIds, numRows); + try (ProjectionOperator projectionOperator = + ProjectionOperatorUtils.getProjectionOperator(_phase2DataSourceMap, docIdOperator, _queryContext)) { + TransformOperator transformOperator = + new TransformOperator(_queryContext, projectionOperator, _nonOrderByExpressions); Review Comment: Agreed that this is a real cost for one-row runs. It is deferred to a follow-up: #19923 (item A1). Reasons: - Results are not affected. Only phase-two setup and chunk decompression are repeated per row. - The fix changes how the leaf builds its blocks. It is planned together with the shared projection helper (review note 2, item A2 in the same issue), because both change the same phase-two code. Note: with the P2a fix, runs are now cut on the whole sorted prefix. More queries can reach this case (for example `ORDER BY day, ts` with both columns sorted and `ts` almost unique). This is recorded in the issue. -- 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]
