rohityadav1993 commented on code in PR #19120: URL: https://github.com/apache/pinot/pull/19120#discussion_r4286521552
########## pinot-core/src/main/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperator.java: ########## @@ -0,0 +1,655 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.core.operator.combine; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Comparator; +import java.util.HashSet; +import java.util.List; +import java.util.PriorityQueue; +import java.util.Set; +import java.util.concurrent.ExecutorService; +import javax.annotation.Nullable; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.common.request.context.OrderByExpressionContext; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.core.common.Operator; +import org.apache.pinot.core.operator.AcquireReleaseColumnsSegmentOperator; +import org.apache.pinot.core.operator.ExplainAttributeBuilder; +import org.apache.pinot.core.operator.blocks.results.BaseResultsBlock; +import org.apache.pinot.core.operator.blocks.results.MetadataResultsBlock; +import org.apache.pinot.core.operator.blocks.results.SelectionResultsBlock; +import org.apache.pinot.core.operator.query.StreamingSelectionOrderByOperator; +import org.apache.pinot.core.operator.streaming.BaseStreamingCombineOperator; +import org.apache.pinot.core.plan.SelectionPlanNode; +import org.apache.pinot.core.query.request.context.QueryContext; +import org.apache.pinot.core.query.selection.SelectionOperatorUtils; +import org.apache.pinot.core.query.utils.OrderByComparatorFactory; +import org.apache.pinot.segment.spi.MutableSegment; +import org.apache.pinot.segment.spi.datasource.DataSourceMetadata; +import org.apache.pinot.spi.exception.QueryErrorCode; +import org.apache.pinot.spi.exception.QueryErrorMessage; +import org.apache.pinot.spi.query.QueryThreadContext; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/// Streaming, lazy combine operator for selection ORDER BY queries whose first order-by expression is an identifier. +/// +/// It performs an incremental k-way heap merge across the per-segment operators in `_operators`, returning +/// globally sorted rows in bounded blocks. Each segment is exposed through a [SegmentCursor] that yields that +/// segment's locally-sorted rows in order: +/// +/// - Segments physically sorted on the first order-by column are backed by +/// [StreamingSelectionOrderByOperator], which is pulled lazily one run/block at a time. +/// - Other (e.g. consuming/unsorted) segments are backed by a single materialized top-K block (any +/// [SelectionResultsBlock]-producing operator such as `SelectionOrderByOperator`); the cursor reads that +/// one block and iterates its rows. +/// +/// A [PriorityQueue] of [SegmentCursor] ordered by the [OrderByComparatorFactory] comparator on +/// each +/// cursor's current head row drives the merge with an at-most-one-head-per-active-segment invariant (the heap holds the +/// cursors themselves, never all rows, which would degenerate into a full heap-sort that materializes everything). Each +/// cycle pops the global-min cursor, appends its head to the current output block, advances that one cursor by a single +/// row, and re-offers it if it still has a head. +/// +/// **Min/max lazy segment activation (pruning).** Cursors are sorted by the first order-by column's min value +/// (ASC) / max value (DESC) reusing the `MinMaxValueContext` idea from +/// [MinMaxValueBasedSelectionOrderByCombineOperator]. A cursor is only activated (its segment acquired and first +/// block read) when the merge frontier reaches its min/max, so once `limit + offset` rows are emitted the +/// remaining segments are never acquired or read. See [#activateEligibleCursors(SegmentCursor)] for the +/// correctness argument. +/// A consuming ([MutableSegment]) segment is always given no min/max and activates up front: its bounds are a snapshot +/// that rows ingested before activation can outgrow, since its plan (and doc count) is only built at activation. +/// With null handling on, a segment whose leading column is not flagged non-null by +/// [SelectionPlanNode#isColumnFlaggedNonNull] (it has nulls, the segment predates the flag, or it is a consuming +/// segment) is given no min/max and always activates. Flagged segments keep their min/max. With +/// null handling off, min/max are used as stored. The flag is read from segment metadata, not the null value vector: +/// that vector is a mapped buffer, and probing it here would acquire segments pruning is about to skip. +/// Its effect shows in the stats: a never-activated segment scans no docs, so `numSegmentsMatched` excludes it while +/// `numSegmentsProcessed` counts every segment. Their difference also counts activated segments that matched no rows. +/// TODO: report the activated count separately; that needs a new `DataTable.MetadataKey`. +/// See https://github.com/apache/pinot/pull/19120#discussion_r3871714063 +/// +/// **Segment acquire/release lifecycle.** A cursor acquires its +/// [AcquireReleaseColumnsSegmentOperator] on activation and releases it only when its child operator is fully +/// drained (acquire-on-activate / release-on-exhaust), rather than per run. This is intentional: the backing +/// [StreamingSelectionOrderByOperator] retains a buffer-backed `ValueBlock` across `nextBlock()` +/// calls in its tail-to-sort mode, so releasing between interleaved runs could read segment buffers after a release +/// under prefetch. Holding the acquire for the cursor's lifetime guarantees no release happens between a cursor's +/// own reads; min/max pruning bounds the number of simultaneously-active (acquired) segments to the merge frontier. +/// The rows handed out by the child operators are already deep-copied to heap `Object[]` (via +/// `RowBasedBlockValueFetcher`), so they remain valid after the segment is released. Any cursors still +/// acquired when the merge ends early (LIMIT reached) or errors out are released via [#releaseAllCursors()]. +/// +/// **Streaming paths only.** This operator is installed only where a +/// [org.apache.pinot.core.query.executor.ResultsBlockStreamer] is present (the +/// MSE leaf, driven by [org.apache.pinot.core.operator.streaming.StreamingInstanceResponseOperator], and the +/// gRPC streaming endpoint). The merge emits many bounded [SelectionResultsBlock]s from successive +/// [#getNextBlock()] calls followed by a final [MetadataResultsBlock]. On the blocking single-stage +/// path the caller expects one complete block from a single [#getNextBlock()] call, so nothing here could +/// stream: the merge would accumulate the whole `limit + offset` result while holding every activated cursor's +/// materialized block alive. That path keeps [MinMaxValueBasedSelectionOrderByCombineOperator]. +/// +/// **Threading.** This operator overrides [#start()]/[#stop()] to no-ops (other than releasing +/// segments) and runs the merge single-threaded and lazily in [#getNextBlock()] on the consumer thread; it does +/// not use the base worker-queue model. The base entry points that it replaces - [#processSegments()] and +/// [#isQuerySatisfied(SelectionResultsBlock, Object)] - are overridden to fail loud. The base +/// `Phaser` (which exists only to fence worker threads against segment release) is intentionally bypassed because +/// all child/segment access is synchronous on the single consumer thread that holds the segment references; no async +/// work may be introduced here without restoring that fence. The instance is single-use (driven once to completion) and +/// is not thread-safe. +@SuppressWarnings({"rawtypes", "unchecked"}) +public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombineOperator<SelectionResultsBlock> { + private static final Logger LOGGER = LoggerFactory.getLogger(StreamingSelectionOrderByCombineOperator.class); + private static final String EXPLAIN_NAME = "COMBINE_SELECT_ORDERBY_STREAMING"; + + private final boolean _asc; + private final boolean _pruningEnabled; + /// Whether a cursor whose bound *ties* the merge frontier may be deferred, not only one sorting strictly + /// beyond it. Only correct under a single order-by expression: see [#sortsBeyond]. + private final boolean _deferTiedCursors; + private final int _numRowsToKeep; + private final int _blockSize; + /// Segments whose metadata marks the leading order-by column sorted; -1 when that expression is not a column + private final int _numSortedSegments; + private final Comparator<Object[]> _comparator; + private final SegmentCursor[] _sortedCursors; + private final PriorityQueue<SegmentCursor> _priorityQueue; + + // Merge progress (single-threaded; mutated only by the consumer thread driving getNextBlock()) + private int _nextToActivate; + private int _numRowsEmitted; + private List<Object[]> _outputRows; + private boolean _done; + /// Captured from the first child block seen; all child blocks share the same schema. Non-null by the time any rows + /// exist to emit: every streaming child emits at least one block carrying its schema, including a zero-match one. + private DataSchema _dataSchema; + /// Deduplicated MERGE_RESPONSE errors for segment blocks dropped on schema mismatch; null until the first mismatch + @Nullable + private Set<String> _dataSchemaMismatchErrors; + /// Subset of the above not yet attached to an emitted block; drained on each attach + @Nullable + private List<String> _unreportedDataSchemaMismatchErrors; + + public StreamingSelectionOrderByCombineOperator(List<Operator> operators, QueryContext queryContext, + ExecutorService executorService) { + // Pass a null merger: we override the consumption path entirely and never touch the base merger / worker queue. + // Every base entry point that would dereference it - processSegments() and isQuerySatisfied() - is overridden to + // throw, so a future change to the base streaming loop fails at that seam instead of with an NPE at query time. + // The executor is unused here (start() is a no-op) but is still handed to the base so that the worker-queue model + // stays constructible; passing null would trade one latent NPE for another. + super(null, operators, queryContext, executorService); + _numRowsToKeep = queryContext.getLimit() + queryContext.getOffset(); + _blockSize = queryContext.getSortedSelectionMergeBlockSize(); + + List<OrderByExpressionContext> orderByExpressions = queryContext.getOrderByExpressions(); + assert orderByExpressions != null && !orderByExpressions.isEmpty(); + OrderByExpressionContext firstOrderByExpression = orderByExpressions.get(0); + _asc = firstOrderByExpression.isAsc(); + _comparator = OrderByComparatorFactory.getComparator(orderByExpressions, queryContext.isNullHandlingEnabled()); + + // Frontier pruning needs a physical column to read min/max from. Under SortedSelectionMergeMode.ON the leading + // order-by may be an expression (forced, for benchmarking), in which case there is no such column: every cursor + // is activated up front in plan order and the merge is a plain k-way heap over materialized children. + ExpressionContext firstOrderByExpressionContext = firstOrderByExpression.getExpression(); + String firstOrderByColumn = + firstOrderByExpressionContext.getType() == ExpressionContext.Type.IDENTIFIER + ? firstOrderByExpressionContext.getIdentifier() : null; + _pruningEnabled = firstOrderByColumn != null; + // A column-0 tie only proves the segment cannot supply an *earlier* row when column 0 is the whole sort key. + // With two or more expressions a tie on column 0 can hide a row sorting earlier on column 1, so ties must + // still activate (same gate as MinMaxValueBasedSelectionOrderByCombineOperator's numOrderByExpressions == 1). + _deferTiedCursors = orderByExpressions.size() == 1; + + // Build one cursor per segment operator and read its first order-by column min/max for lazy activation ordering. + // Reading DataSourceMetadata does not touch column buffers, so no segment acquire is needed here (mirrors + // MinMaxValueBasedSelectionOrderByCombineOperator). + _sortedCursors = new SegmentCursor[_numOperators]; + int numSortedSegments = firstOrderByColumn != null ? 0 : -1; + for (int i = 0; i < _numOperators; i++) { + Operator<BaseResultsBlock> operator = _operators.get(i); + if (firstOrderByColumn == null) { + _sortedCursors[i] = new SegmentCursor(operator, null, null); + continue; + } + DataSourceMetadata metadata = + operator.getIndexSegment().getDataSource(firstOrderByColumn, queryContext.getSchema()) + .getDataSourceMetadata(); + Comparable minValue = metadata.getMinValue(); + Comparable maxValue = metadata.getMaxValue(); + // A consuming segment's min/max is a snapshot copied when getDataSource() builds the MutableDataSource, but under + // prefetch (AcquireReleaseColumnsSegmentOperator) the child plan, and with it the doc count, is only built at + // activation. Rows indexed in between can fall outside these bounds and be emitted out of order, so treat the + // bounds as unknown and let activateEligibleCursors force-activate this cursor. + if (operator.getIndexSegment() instanceof MutableSegment) { + minValue = null; + maxValue = null; + } + // isNonNull is loaded with the segment, the same class of access as min/max. False means nulls or unknown, so + // drop the bounds and let activateEligibleCursors force-activate this cursor. Do not read the null vector here. + if (queryContext.isNullHandlingEnabled() + && !SelectionPlanNode.isColumnFlaggedNonNull(operator.getIndexSegment(), firstOrderByColumn)) { + minValue = null; + maxValue = null; + } + _sortedCursors[i] = new SegmentCursor(operator, minValue, maxValue); + if (metadata.isSorted()) { + numSortedSegments++; + } + } + _numSortedSegments = numSortedSegments; + if (firstOrderByColumn != null) { + sortCursorsByMinMax(); + } + + _priorityQueue = new PriorityQueue<>(Math.max(1, _numOperators), + (o1, o2) -> _comparator.compare(o1.currentHead(), o2.currentHead())); + _outputRows = newOutputList(); + } + + /// Sorts the cursors so the merge can activate them lazily in frontier order: ascending by the column min value for + /// ASC, descending by the column max value for DESC. Cursors without a min/max are placed first because they must + /// always be processed (mirrors [MinMaxValueBasedSelectionOrderByCombineOperator]). + private void sortCursorsByMinMax() { + if (_asc) { + Arrays.sort(_sortedCursors, (o1, o2) -> { + if (o1._minValue == null) { + return o2._minValue == null ? 0 : -1; + } + if (o2._minValue == null) { + return 1; + } + return o1._minValue.compareTo(o2._minValue); + }); + } else { + Arrays.sort(_sortedCursors, (o1, o2) -> { + if (o1._maxValue == null) { + return o2._maxValue == null ? 0 : -1; + } + if (o2._maxValue == null) { + return 1; + } + return o2._maxValue.compareTo(o1._maxValue); + }); + } + } + + @Override + public String toExplainString() { + return EXPLAIN_NAME; + } + + /// Settings fixed at construction. Explain never drives the merge, so how many segments activate is not shown here. + /// Query-wide values are idempotent so the broker merges the node across workers only when they agree; the segment + /// counts are summed. `numSortedSegments` is metadata sortedness, not whether a child streams: a DESC scan without + /// `allowReverseOrder`, or a leading column holding nulls, still falls back to a materialized child. + @Override + protected void explainAttributes(ExplainAttributeBuilder attributeBuilder) { + super.explainAttributes(attributeBuilder); + attributeBuilder.putLongIdempotent("blockSize", _blockSize); + attributeBuilder.putBool("frontierPruning", _pruningEnabled); + attributeBuilder.putBool("deferTiedCursors", _deferTiedCursors); + attributeBuilder.putLong("numSegments", _numOperators); + if (_numSortedSegments >= 0) { + attributeBuilder.putLong("numSortedSegments", _numSortedSegments); + } + } + + /// Override to a no-op: the merge is single-threaded and lazy in [#getNextBlock()], so we do not spin up the + /// base worker threads / blocking-queue model. + @Override + public void start() { + } + + /// Override the base worker-queue stop: no worker threads / phaser tasks were started. Release any segments still + /// acquired (idempotent) so an early stop by the driver cannot leak acquires. + @Override + public void stop() { + _done = true; + releaseAllCursors(); + } + + /// The base worker-thread entry point must never run here ([#start()] is a no-op). Fail loud if it ever does. + @Override + protected void processSegments() { + throw new IllegalStateException( + "StreamingSelectionOrderByCombineOperator runs single-threaded; processSegments() must not be called"); + } + + /// Only reachable from the base streaming loop, which this operator replaces; the base implementation would + /// dereference the null merger. Fail loud rather than NPE if a future change routes back through it. + @Override + protected boolean isQuerySatisfied(SelectionResultsBlock resultsBlock, Object tracker) { + throw new IllegalStateException( + "StreamingSelectionOrderByCombineOperator tracks its own limit; isQuerySatisfied() must not be called"); + } + + @Override + protected BaseResultsBlock getNextBlock() { + if (_done) { + // Streaming mode: terminal metadata block after the last data block. Idempotent if called again. + return attachExecutionStats(new MetadataResultsBlock()); Review Comment: Done in 1ef67c18b6, in InstanceResponseOperator.getBaseBlock(): main-thread CPU, allocated bytes and wall clock are now summed across calls, and the worker totals are still taken from the latest block since they already cover the whole query. The blocking path makes one call and is unchanged. This also fixes the smaller undercount for the other streaming combines (their consumer-side polling time). Test: the real combine through StreamingInstanceResponseOperator over about 2,000 streamed blocks, where each child block allocates a known amount on the consumer thread; the response must report at least that much memory and CPU. Before the fix it reported about 24 KB and 1 ms against at least 132 MB and 72 ms. One visible change: system-activities CPU for a streaming response is now derived from the summed in-operator wall clock, which excludes time blocked in send. For this combine it now reflects real off-CPU time instead of about 0. -- 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]
