You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by xi...@apache.org on 2026/09/30 15:54:04 UTC
(pinot) branch xiangfu0/codex1/sorted-exchange-leaf updated (42e3b467ec6 -> eecbbd1c044)
This is an automated email from the ASF dual-hosted git repository.
xiangfu0 pushed a change to branch xiangfu0/codex1/sorted-exchange-leaf
in repository https://gitbox.apache.org/repos/asf/pinot.git
omit 42e3b467ec6 Check the DESC null-bearing case against fixed rows instead of MinMax
omit 11b0adc18fc Keep frontier pruning for segments flagged non-null.
omit be55e91d1a4 Address second-round review on the streaming selection ORDER BY combine
omit 8de76d68cd3 Defer tied cursors under a single-column ORDER BY
omit fcde3855208 Address review feedback on the streaming selection ORDER BY combine
omit 590e70ffb18 Add streaming selection ORDER BY combine for physically sorted segments
omit 5dc4fb34dda Check cancellation while validating merged ordering
omit 35564e66b37 Stream ORDER BY over confirmed sorted mailbox producers
omit 68436765f71 [MSE] Speed up four-sender sorted window merge
omit d421c401f2a [MSE] Address sender-sorted window review feedback
omit 3f45fcef683 [MSE] Minimize sender-sorted merge changes
omit 291e1d106f4 [MSE] Avoid eager sort buffer allocation
omit 4d586bf499b [MSE] Limit sender-sorted merge to global windows
omit 9e9de1a1a45 [MSE] Optimize sorted mailbox merge
omit 299551d6df5 [MSE] Merge sender-sorted streams for ordered windows
add cd504afe901 Add JDBC support for MAP, arrays and extended scalar types
add f5614058a96 Bump com.h2database:h2 from 2.5.250 to 2.5.252 (#19719)
add a42e748472a Bump brace-expansion in /pinot-controller/src/main/resources (#19721)
add f835005ae52 [MSE] Merge sender-sorted streams for ordered windows
add a52f8a91c40 [MSE] Optimize sorted mailbox merge
add 63b3a30cb15 [MSE] Limit sender-sorted merge to global windows
add a96db4db195 [MSE] Avoid eager sort buffer allocation
add f888a8d946d [MSE] Minimize sender-sorted merge changes
add 36ccd853f40 [MSE] Address sender-sorted window review feedback
add a3ac0c848eb [MSE] Speed up four-sender sorted window merge
add d9dcc8687e4 [MSE] Simplify sorted mailbox merging
add 555f55c106e [MSE] Use current planner test APIs
add 1374926823f Stream ORDER BY over confirmed sorted mailbox producers
add 4d1339a3bae Check cancellation while validating merged ordering
add 023b8d73095 Simplify ordered mailbox test fixtures
add 697a7106b9d [MSE] Preserve merge test imports after restacking
new 849901a015d Add streaming selection ORDER BY combine for physically sorted segments
new 17f1c3aa645 Address review feedback on the streaming selection ORDER BY combine
new 0da93353bfc Defer tied cursors under a single-column ORDER BY
new a8065e5459d Address second-round review on the streaming selection ORDER BY combine
new 7efe592da81 Keep frontier pruning for segments flagged non-null.
new 164da0d2252 Check the DESC null-bearing case against fixed rows instead of MinMax
new eecbbd1c044 Share sorted selection setup and trim duplicate coverage
This update added new revisions after undoing existing revisions.
That is to say, some revisions that were in the old version of the
branch are not in the new version. This situation occurs
when a user --force pushes a change and generates a repository
containing something like this:
* -- * -- B -- O -- O -- O (42e3b467ec6)
\
N -- N -- N refs/heads/xiangfu0/codex1/sorted-exchange-leaf (eecbbd1c044)
You should already have received notification emails for all of the O
revisions, and so the following emails describe only the N revisions
from the common base, B.
Any revisions marked "omit" are not gone; other references still
refer to them. Any revisions marked "discard" are gone forever.
The 7 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
.../org/apache/pinot/client/PinotResultSet.java | 281 +++++-----------
.../pinot/client/base/AbstractBaseResultSet.java | 342 ++++++++++++++++++-
.../pinot/client/grpc/PinotGrpcResultSet.java | 293 ++++++-----------
.../apache/pinot/client/utils/BigDecimalUtils.java | 20 +-
.../apache/pinot/client/utils/DateTimeUtils.java | 16 +-
.../org/apache/pinot/client/utils/DriverUtils.java | 20 ++
.../apache/pinot/client/PinotResultSetTest.java | 164 +++++++--
.../pinot/client/grpc/PinotGrpcResultSetTest.java | 260 +++++++++++++++
.../pinot/client/utils/BigDecimalUtilsTest.java | 35 +-
.../response/encoder/JsonResponseEncoder.java | 95 +-----
.../org/apache/pinot/common/utils/DataSchema.java | 19 ++
.../src/main/resources/package-lock.json | 6 +-
.../StreamingSelectionOrderByCombineOperator.java | 32 +-
.../query/StreamingSelectionOrderByOperator.java | 121 +++----
...reamingSelectionOrderByCombineOperatorTest.java | 74 +----
.../StreamingSelectionOrderByOperatorTest.java | 342 ++-----------------
.../apache/pinot/query/QueryCompilationTest.java | 18 +-
.../query/planner/logical/PlanFragmenterTest.java | 27 +-
.../SortedMailboxMergeReceiveOperator.java | 366 +--------------------
.../query/runtime/operator/LeafOperatorTest.java | 18 +-
.../query/runtime/operator/SortOperatorTest.java | 9 +-
.../SortedMailboxMergeReceiveOperatorTest.java | 110 -------
.../plan/OpChainConverterDispatcherTest.java | 7 +-
.../plan/server/ServerPlanRequestVisitorTest.java | 18 +-
pom.xml | 2 +-
25 files changed, 1184 insertions(+), 1511 deletions(-)
copy pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/CancelObserver.java => pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/utils/BigDecimalUtils.java (65%)
create mode 100644 pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/grpc/PinotGrpcResultSetTest.java
copy pinot-spi/src/test/java/org/apache/pinot/spi/utils/builder/TableConfigBuilderTest.java => pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/utils/BigDecimalUtilsTest.java (52%)
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
(pinot) 03/07: Defer tied cursors under a single-column ORDER BY
Posted by xi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
xiangfu0 pushed a commit to branch xiangfu0/codex1/sorted-exchange-leaf
in repository https://gitbox.apache.org/repos/asf/pinot.git
commit 0da93353bfcf305806218c14e6bb241ea52fcb90
Author: rohity <ro...@uber.com>
AuthorDate: Sat Sep 26 18:38:51 2026 +0000
Defer tied cursors under a single-column ORDER BY
Surfaced by the arm-4 memory benchmark, not by review. Where segment
min/max values tie on the leading ORDER BY column -- the normal shape
for a low-cardinality or timestamp-prefix sort key -- every cursor
activated at once, each pinning a decompressed block, because
sortsBeyond() only deferred a cursor sorting strictly past the merge
frontier. MinMaxValueBasedSelectionOrderByCombineOperator has had the
tie case since it was written; this path did not.
Relax the comparison to non-strict, gated on a single order-by
expression. With two or more, a column-0 tie can hide a row sorting
earlier on column 1, which would be genuine out-of-order emission --
the same gate the MinMax operator applies for the same reason.
The justification differs from that operator's, though, and the javadoc
says so: its bound is the k-th row of a complete top-K, so a tie
provably cannot improve the answer. Here the bound is the live merge
frontier, taken while fewer than limit + offset rows have been emitted,
so that argument is unavailable. What holds instead is that this only
ever defers: _nextToActivate does not advance on a deferral, and the
caller force-activates once the heap and leader both drain. Since the
bound bounds every row in the segment, a deferred cursor holds no row
sorting strictly before one already emitted -- only rows tying it,
interchangeable when column 0 is the whole sort key.
The frontier is the smaller of the retained leader's head and the heap
top's head, since that is the next row emitted. Defer when the bound
sorts past either head, which is exactly past the smaller one; checking
against the larger overstated the frontier whenever the leader had moved
past the heap head, and every remaining tied cursor activated. A present
candidate with a null head still forces activation.
This does change which tied rows a single-column ORDER BY returns. The
class already declares tie order arbitrary for rows equal on every
order-by expression, which under one column is the same set.
Tests on the tied-minima fixture: the deferral itself ASC and DESC
(both fail without the change, scanning every segment instead of one),
value-multiset correctness across a LIMIT straddling a tied group and
across an OFFSET, the two-expression gate holding full row parity, no
under-delivery when the LIMIT covers every row, and one segment
activated per tied run once the leader moves past the heap head, ASC,
DESC, and across output blocks.
---
.../StreamingSelectionOrderByCombineOperator.java | 50 ++++--
...reamingSelectionOrderByCombineOperatorTest.java | 171 +++++++++++++++++++++
2 files changed, 208 insertions(+), 13 deletions(-)
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperator.java
index 527e7a77fb7..75560d3715b 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperator.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperator.java
@@ -110,6 +110,9 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
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 {@link #sortsBeyond}.
+ private final boolean _deferTiedCursors;
private final int _numRowsToKeep;
private final int _blockSize;
private final Comparator<Object[]> _comparator;
@@ -156,6 +159,10 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
firstOrderByExpressionContext.getType() == ExpressionContext.Type.IDENTIFIER
? firstOrderByExpressionContext.getIdentifier() : null;
_pruningEnabled = firstOrderByColumn != null && !queryContext.isNullHandlingEnabled();
+ // 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
@@ -325,11 +332,15 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
/// Cursors are visited in min/max order and {@code _nextToActivate} advances only on a real activation, so a
/// {@code break} defers the current cursor to a later call with a risen frontier rather than skipping it. Deferring
/// is safe because the first order-by column is the primary sort key: a cursor whose range starts past the frontier
- /// holds no row sorting before it. With no frontier known yet, activation is forced.
+ /// holds no row sorting before it. With no frontier known yet, activation is forced. Under a single order-by
+ /// expression a cursor whose bound merely *ties* the frontier is deferred too, which is what keeps a long run of
+ /// equal segment minima from activating every segment at once; see {@link #sortsBeyond}.
///
/// The frontier is the row about to be emitted. Understating it defers a cursor that could have supplied that row
/// and the merge emits out of order, where overstating it only activates a segment early. The leader is retained
- /// outside the heap, so that row is either its head or the heap head: a bound past both is past the frontier.
+ /// outside the heap, so that row is the smaller of its head and the heap head: a bound past either is past the
+ /// frontier. Checking the leader alone would overstate it whenever the leader has moved past the heap head, which
+ /// happens both on entry (the leader just advanced) and within this loop (an activation offers a smaller head).
private void activateEligibleCursors(@Nullable SegmentCursor leader) {
while (_nextToActivate < _sortedCursors.length) {
SegmentCursor cursor = _sortedCursors[_nextToActivate];
@@ -338,7 +349,7 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
// Re-read per iteration: each activation below can offer a smaller head into the heap.
SegmentCursor top = _priorityQueue.peek();
// A null bound always activates, and with no candidate at all there is no frontier to prune against.
- if (bound != null && (leader != null || top != null) && sortsBeyond(bound, leader) && sortsBeyond(bound, top)) {
+ if (bound != null && (leader != null || top != null) && sortsBeyond(bound, leader, top)) {
break;
}
}
@@ -350,21 +361,34 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
}
}
- /// Whether `bound` sorts past the cursor's head on the first order-by column, so that segment cannot supply the
- /// head. An absent cursor imposes no constraint; a null head cannot be compared, so this reports `false` and the
- /// caller activates. Tests column 0 only, as the pruning bound always has, rather than the full-row comparator --
- /// which would put every order-by column on the per-row path.
- private boolean sortsBeyond(Comparable bound, @Nullable SegmentCursor cursor) {
- if (cursor == null) {
- return true;
- }
- Object headValue = cursor.currentHead()[0];
- if (headValue == null) {
+ /// Whether `bound` sorts past the frontier -- the smaller of the two candidates' heads -- on the first order-by
+ /// column, so that segment cannot supply the next row. An absent candidate imposes no constraint; a null head cannot
+ /// be compared, so this reports `false` and the caller activates. Tests column 0 only, as the pruning bound always
+ /// has, rather than the full-row comparator -- which would put every order-by column on the per-row path.
+ ///
+ /// Under a single order-by expression ({@code _deferTiedCursors}) a *tie* also defers, which only postpones a
+ /// cursor: {@code _nextToActivate} does not advance, and the caller force-activates once heap and leader drain.
+ /// Since {@code bound} bounds every row in the segment, a deferred cursor holds no row sorting strictly before an
+ /// emitted one -- only ties, interchangeable when column 0 is the whole sort key. Note this is *not*
+ /// {@link MinMaxValueBasedSelectionOrderByCombineOperator}'s justification: its bound is a complete top-K's k-th
+ /// row, this one is the live frontier. It changes which tied rows are returned, already documented as arbitrary.
+ private boolean sortsBeyond(Comparable bound, @Nullable SegmentCursor leader, @Nullable SegmentCursor top) {
+ Object leaderHead = leader != null ? leader.currentHead()[0] : null;
+ Object topHead = top != null ? top.currentHead()[0] : null;
+ if ((leader != null && leaderHead == null) || (top != null && topHead == null)) {
return false;
}
+ // Past either head is past the smaller one.
+ return (leader != null && sortsBeyond(bound, leaderHead)) || (top != null && sortsBeyond(bound, topHead));
+ }
+
+ private boolean sortsBeyond(Comparable bound, Object headValue) {
// Both come from the same first order-by column: the metadata min/max and the materialized row[0] share the
// column's stored type, so this comparison is type-safe (same assumption as MinMaxValueBased...).
int cmp = bound.compareTo(headValue);
+ if (_deferTiedCursors) {
+ return _asc ? cmp >= 0 : cmp <= 0;
+ }
return _asc ? cmp > 0 : cmp < 0;
}
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java
index 0ba5d5eb5bf..ae269cc1e76 100644
--- a/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java
+++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java
@@ -390,6 +390,167 @@ public class StreamingSelectionOrderByCombineOperatorTest {
"SELECT sortedCol, tailCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 40", false);
}
+ /// Pins the tie-deferral itself: under a single order-by expression, every `_lowCardSegments` segment's sortedCol
+ /// minimum ties at 0, so a LIMIT satisfiable from one segment's leading run of 0s (25 rows) must not activate any
+ /// other segment. Before `_deferTiedCursors` existed, every cursor tied the (unknown-yet) frontier at activation
+ /// time and all of them activated up front, each scanning its first block (about LIMIT docs), so docs scanned
+ /// pre-fix is about `NUM_SEGMENTS * 20`. Post-fix only the leading segment (in cursor order) is ever touched, so docs
+ /// scanned is bounded by one segment's worth.
+ @Test
+ public void testTieDeferralBoundsDocsScannedToLeadingSegment() {
+ Result result = run(_lowCardSegments, "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol LIMIT 20", true,
+ false, true, 0);
+ assertTrue(result._combineOperator instanceof StreamingSelectionOrderByCombineOperator);
+ assertEquals(result._rows.size(), 20);
+ for (Object[] row : result._rows) {
+ assertEquals((int) row[0], 0, "LIMIT 20 must be satisfiable entirely from sortedCol=0 rows");
+ }
+ assertTrue(result._numDocsScanned <= NUM_RECORDS_PER_SEGMENT,
+ "A LIMIT satisfiable from one segment's leading run must not activate any other tied segment; docs scanned: "
+ + result._numDocsScanned);
+ }
+
+ /// DESC counterpart of [#testTieDeferralBoundsDocsScannedToLeadingSegment]: every `_lowCardSegments` segment's
+ /// sortedCol maximum ties at 3, so `sortsBeyond`'s `_asc ? cmp >= 0 : cmp <= 0` DESC branch is the one under test
+ /// here rather than the ASC one above. A sign flip in that branch would either over-scan (fail this test the same
+ /// way the pre-fix code fails the ASC test) or, worse, prune a segment that still had rows to give -- which the
+ /// row-count / row-value assertions below would catch.
+ @Test
+ public void testTieDeferralBoundsDocsScannedToLeadingSegmentDesc() {
+ Result result = run(_lowCardSegments,
+ "SET allowReverseOrder=true; SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol DESC LIMIT 20", true,
+ false, true, 0);
+ assertTrue(result._combineOperator instanceof StreamingSelectionOrderByCombineOperator);
+ assertEquals(result._rows.size(), 20);
+ for (Object[] row : result._rows) {
+ assertEquals((int) row[0], 3, "LIMIT 20 DESC must be satisfiable entirely from sortedCol=3 rows");
+ }
+ assertTrue(result._numDocsScanned <= NUM_RECORDS_PER_SEGMENT,
+ "A LIMIT satisfiable from one segment's leading run must not activate any other tied segment; docs scanned: "
+ + result._numDocsScanned);
+ }
+
+ /// Correctness under single-column ties, LIMIT straddling a tied group: with `ORDER BY sortedCol` alone every
+ /// segment ties at every one of the 4 distinct values, so a LIMIT of 30 forces the merge past the sortedCol=0
+ /// boundary (only 25 such rows per segment) partway through a block from the leading segment. A wrongly-timed
+ /// reactivation here would either drop true sortedCol=0 rows from other segments (undercount at the tie) or emit
+ /// out of order. [#assertParity]'s full-row multiset check does not apply: with sortedCol as the only key, a
+ /// different (equally valid) tied row may be returned than the MinMax baseline picks, so only the order-by column
+ /// values -- whose multiset *is* pinned by the LIMIT boundary, unlike the underlying rows -- are compared.
+ @Test
+ public void testSingleColumnTieDeferralPreservesOrderByValuesAcrossLimitStraddle() {
+ @Language("sql") String query = "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol LIMIT 30";
+ Result baseline = run(_lowCardSegments, query, false, false, false, 0);
+ Result streamed = run(_lowCardSegments, query, true, false, true, 3);
+ assertTrue(streamed._combineOperator instanceof StreamingSelectionOrderByCombineOperator);
+ assertEquals(streamed._rows.size(), 30);
+ assertSorted(streamed._rows, orderByComparator(query, false));
+ assertEquals(orderByColumnValues(streamed._rows), orderByColumnValues(baseline._rows),
+ "Multiset of sortedCol values must match the MinMax baseline even though the underlying rows may differ");
+ }
+
+ /// Correctness under single-column ties with an OFFSET: LIMIT 10 OFFSET 20 straddles the same sortedCol=0/1
+ /// boundary as [#testSingleColumnTieDeferralPreservesOrderByValuesAcrossLimitStraddle] (limit + offset = 30), but
+ /// exercises it with a nonzero offset. As [#testLimitOffsetParity] documents, the server (and this combine) retains
+ /// {@code limit + offset} rows -- the offset is trimmed by the broker afterwards -- so 30, not 10, rows come back
+ /// here too; what differs from the straddle test is only the query shape, not the row count.
+ @Test
+ public void testSingleColumnTieDeferralPreservesOrderByValuesAcrossOffset() {
+ @Language("sql") String query = "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol LIMIT 10 OFFSET 20";
+ Result baseline = run(_lowCardSegments, query, false, false, false, 0);
+ Result streamed = run(_lowCardSegments, query, true, false, true, 3);
+ assertTrue(streamed._combineOperator instanceof StreamingSelectionOrderByCombineOperator);
+ assertEquals(streamed._rows.size(), 30, "Server retains limit + offset rows; the broker trims the offset later");
+ assertSorted(streamed._rows, orderByComparator(query, false));
+ assertEquals(orderByColumnValues(streamed._rows), orderByColumnValues(baseline._rows),
+ "Multiset of sortedCol values must match the MinMax baseline even though the underlying rows may differ");
+ }
+
+ /// Condition 1 (the two-expression gate): `_lowCardSegments` still tie on sortedCol alone, but a second order-by
+ /// expression (valCol) makes the full order-by key a total order, so full-row parity applies unlike the
+ /// single-column tests above. `_deferTiedCursors` must be false here (`orderByExpressions.size() == 1` fails), so
+ /// this pins the same shape both before and after the production change: a wrongly-deferred cursor on this shape
+ /// would drop a row that sorts earlier on valCol despite tying on sortedCol, which parity would catch as a missing
+ /// row rather than merely a reordered one.
+ @Test
+ public void testTwoExpressionOrderByDoesNotDeferTiedCursors() {
+ assertParity(_lowCardSegments, "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 30",
+ false);
+ }
+
+ /// Condition 2 (no under-delivery when the heap drains): a single-column ORDER BY with a LIMIT covering every row
+ /// of `_lowCardSegments` forces every segment to eventually activate no matter how aggressively ties are deferred
+ /// -- deferral only postpones a cursor, it never removes it from consideration, and the
+ /// `(leader != null || top != null)` guard in `activateEligibleCursors` forces activation once the heap and leader
+ /// both go empty. Asserts the full row count comes back with none missing.
+ @Test
+ public void testSingleColumnTieDeferralNeverUnderDeliversAtFullDrain() {
+ int totalRows = NUM_SEGMENTS * NUM_RECORDS_PER_SEGMENT;
+ @Language("sql") String query =
+ "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol LIMIT " + totalRows;
+ Result baseline = run(_lowCardSegments, query, false, false, false, 0);
+ Result streamed = run(_lowCardSegments, query, true, false, true, 7);
+ assertTrue(streamed._combineOperator instanceof StreamingSelectionOrderByCombineOperator);
+ assertEquals(streamed._rows.size(), totalRows, "Every row must be delivered when the LIMIT covers the whole set");
+ assertSorted(streamed._rows, orderByComparator(query, false));
+ assertEquals(orderByColumnValues(streamed._rows), orderByColumnValues(baseline._rows),
+ "Multiset of sortedCol values must match the MinMax baseline when the merge drains completely");
+ }
+
+ /// Each `_lowCardSegments` segment holds only 25 rows at the tied minimum, so a LIMIT past 25 moves the leader's head
+ /// to sortedCol=1 while the other segments still tie at 0. The next row out is then the smaller of the leader's and
+ /// the heap's heads, and a waiting cursor must be deferred against that row, not against the leader alone: checking
+ /// the leader activates every remaining tied segment the moment one of them has been opened. Activating one segment
+ /// per 25-row run is the minimum, so at most `ceil(limit / 25)` segments are matched (a segment is matched once
+ /// activation reads its first block).
+ @Test
+ public void testTieDeferralActivatesOneSegmentPerTiedRunPastLeaderRun() {
+ for (int limit : new int[]{30, 60, 90}) {
+ Result result = run(_lowCardSegments, "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol LIMIT " + limit,
+ true, false, true, 0);
+ assertTrue(result._combineOperator instanceof StreamingSelectionOrderByCombineOperator);
+ assertEquals(result._rows.size(), limit);
+ int numZeros = Math.min(limit, 25 * NUM_SEGMENTS);
+ assertEquals(result._rows.stream().filter(row -> (int) row[0] == 0).count(), numZeros,
+ "Every returned row must come from the tied minimum while enough such rows exist");
+ int maxSegments = (limit + 24) / 25;
+ assertTrue(result._numSegmentsMatched <= maxSegments,
+ "LIMIT " + limit + " should activate at most " + maxSegments + " segments; activated: "
+ + result._numSegmentsMatched);
+ }
+ }
+
+ /// DESC counterpart of [#testTieDeferralActivatesOneSegmentPerTiedRunPastLeaderRun]: the leader leaves the tied
+ /// maximum (3) for 2 while the waiting segments still tie at 3.
+ @Test
+ public void testTieDeferralActivatesOneSegmentPerTiedRunPastLeaderRunDesc() {
+ Result result = run(_lowCardSegments,
+ "SET allowReverseOrder=true; SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol DESC LIMIT 30", true,
+ false, true, 0);
+ assertTrue(result._combineOperator instanceof StreamingSelectionOrderByCombineOperator);
+ assertEquals(result._rows.size(), 30);
+ for (Object[] row : result._rows) {
+ assertEquals((int) row[0], 3, "LIMIT 30 DESC must be satisfiable entirely from sortedCol=3 rows");
+ }
+ assertTrue(result._numSegmentsMatched <= 2,
+ "LIMIT 30 DESC should activate at most 2 segments; activated: " + result._numSegmentsMatched);
+ }
+
+ /// Same shape with a 7-row output block, so the leader is handed back to the heap and re-chosen mid-run: deferral
+ /// against the next row must hold across block boundaries too, and the rows must still match the baseline's values.
+ @Test
+ public void testTieDeferralPastLeaderRunAcrossOutputBlocks() {
+ @Language("sql") String query = "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol LIMIT 60";
+ Result baseline = run(_lowCardSegments, query, false, false, false, 0);
+ Result streamed = run(_lowCardSegments, query, true, false, true, 7);
+ assertTrue(streamed._combineOperator instanceof StreamingSelectionOrderByCombineOperator);
+ assertEquals(streamed._rows.size(), 60);
+ assertSorted(streamed._rows, orderByComparator(query, false));
+ assertEquals(orderByColumnValues(streamed._rows), orderByColumnValues(baseline._rows));
+ assertTrue(streamed._numSegmentsMatched <= 3,
+ "LIMIT 60 should activate at most 3 segments; activated: " + streamed._numSegmentsMatched);
+ }
+
@Test
public void testTwoPhaseSelectNonOrderByParity() {
// tailCol is selected but not an order-by key -> the streaming children take the two-phase (order-by-then-fetch)
@@ -964,6 +1125,7 @@ public class StreamingSelectionOrderByCombineOperatorTest {
result._schema = block.getDataSchema();
}
result._numDocsScanned = block.getNumDocsScanned();
+ result._numSegmentsMatched = block.getNumSegmentsMatched();
break;
}
SelectionResultsBlock dataBlock = (SelectionResultsBlock) block;
@@ -1026,6 +1188,14 @@ public class StreamingSelectionOrderByCombineOperatorTest {
}).sorted().collect(Collectors.toList());
}
+ /// Extracts the ORDER BY column (projected at index 0 in every query these tests use) as a sorted multiset,
+ /// ignoring every other projected column. Used where a full-row [#assertMultisetEquals] would be too strong: under
+ /// a single order-by expression a tied value may be satisfied by different underlying rows than the MinMax
+ /// baseline picks, but the count of each order-by value in the result is still pinned by the LIMIT boundary.
+ private static List<Integer> orderByColumnValues(List<Object[]> rows) {
+ return rows.stream().map(row -> (Integer) row[0]).sorted().collect(Collectors.toList());
+ }
+
@AfterClass
public void tearDown()
throws IOException {
@@ -1047,5 +1217,6 @@ public class StreamingSelectionOrderByCombineOperatorTest {
private List<Integer> _blockSizes;
private int _numBlocks;
private long _numDocsScanned;
+ private int _numSegmentsMatched;
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
(pinot) 02/07: Address review feedback on the streaming selection ORDER BY combine
Posted by xi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
xiangfu0 pushed a commit to branch xiangfu0/codex1/sorted-exchange-leaf
in repository https://gitbox.apache.org/repos/asf/pinot.git
commit 17f1c3aa64506e3ccccda4c2deef5798d1dafa96
Author: rohity <ro...@uber.com>
AuthorDate: Thu Sep 17 18:01:08 2026 +0000
Address review feedback on the streaming selection ORDER BY combine
Eight rounds of review on apache/pinot#19120, folded into one commit.
The streaming merge stays opt-in and off by default, so a query that
does not ask for it is unaffected by everything below.
- Restrict the path to plans that have a ResultsBlockStreamer. On the
blocking path the caller takes a single nextBlock() and expects the
whole result, so the merge could never flush and would hold every
activated cursor's block alive for nothing.
- Replace the boolean streamingSelectionOrderBy option with the
tri-valued sortedSelectionMergeMode (OFF / ON / AUTO, default OFF;
the block size option is renamed to match), and resolve AUTO once in
InstancePlanMakerImplV2 rather than at the gates. CombinePlanNode
builds its leaves before evaluating its own gate, so a decision there
arrives too late to stop streaming leaves being built under a
non-streaming combine. AUTO takes the merge only when at least
sortedSelectionMergeAutoMinSortedRatio of the queried segments are
physically sorted on the leading ORDER BY column, decided from
segment metadata alone before any segment is acquired.
- Resolve AUTO to OFF for a DESC leading expression without
allowReverseOrder. DataSourceMetadata.isSorted() reports ascending
physical order, so such a query cannot stream, and taking the path
would give up the MinMax combine's parallelism and min/max pruning
in exchange for nothing.
- Route both scan paths through one empty-block skip. The two paths
disagreed about the project operator's contract: the sorted path
read an empty block as exhaustion and would have ended a scan
mid-segment.
- Emit one schema-carrying empty block from a segment that matches no
rows, so a consumer can always learn the child's DataSchema, and
derive the two-phase schema from column metadata instead of standing
up a projection and transform operator per empty segment.
- Hold the winning cursor outside the heap. Segments are near-disjoint
on a time-like leading column, so one cursor supplies long runs of
consecutive rows, each of which previously paid two O(log k) sift
operations to arrive back at the same cursor.
- Return a growable row list from the tail-to-sort drain, and fail
loud on the two base worker entry points the combine replaces, so a
future change routing back through them breaks at the seam instead
of dereferencing the null merger at query time.
- Reuse the sorted-by project on the DESC fall-through instead of
building a second identical one.
Every behavioural change above carries a regression test verified to
fail against the unfixed code.
---
.../common/utils/config/QueryOptionsUtils.java | 48 +-
.../common/utils/config/QueryOptionsUtilsTest.java | 49 ++
.../StreamingSelectionOrderByCombineOperator.java | 251 +++++----
.../query/StreamingSelectionOrderByOperator.java | 149 ++++--
.../apache/pinot/core/plan/CombinePlanNode.java | 22 +-
.../apache/pinot/core/plan/SelectionPlanNode.java | 75 ++-
.../core/plan/maker/InstancePlanMakerImplV2.java | 111 +++-
.../core/query/request/context/QueryContext.java | 54 +-
.../operator/combine/CombineSlowOperatorsTest.java | 25 +-
...reamingSelectionOrderByCombineOperatorTest.java | 588 +++++++++++++++++++--
.../StreamingSelectionOrderByOperatorTest.java | 471 ++++++++++++++++-
.../apache/pinot/spi/utils/CommonConstants.java | 43 +-
12 files changed, 1626 insertions(+), 260 deletions(-)
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java
index d8393720613..06dd7b78b6a 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java
@@ -39,6 +39,7 @@ import org.apache.pinot.spi.utils.CommonConstants;
import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey;
import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner.JoinOverFlowMode;
import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner.WindowOverFlowMode;
+import org.apache.pinot.spi.utils.CommonConstants.Server.SortedSelectionMergeMode;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -602,8 +603,28 @@ public class QueryOptionsUtils {
return checkedParseIntNonNegative(QueryOptionKey.STREAMING_DISTINCT_FLUSH_THRESHOLD, value);
}
- public static boolean isSortedSelectionMergeEnabled(Map<String, String> queryOptions) {
- return Boolean.parseBoolean(queryOptions.get(QueryOptionKey.SORTED_SELECTION_MERGE_ENABLED));
+ /// Returns the streaming selection ORDER BY merge mode, or null if not set. Accepts the enum constant names in
+ /// any case (`on` / `off` / `auto`).
+ @Nullable
+ public static SortedSelectionMergeMode getSortedSelectionMergeMode(Map<String, String> queryOptions) {
+ String value = queryOptions.get(QueryOptionKey.SORTED_SELECTION_MERGE_MODE);
+ if (value == null) {
+ return null;
+ }
+ try {
+ return SortedSelectionMergeMode.valueOf(value.trim().toUpperCase());
+ } catch (IllegalArgumentException e) {
+ throw new IllegalArgumentException(
+ String.format("%s must be one of %s, got: %s", QueryOptionKey.SORTED_SELECTION_MERGE_MODE,
+ Arrays.toString(SortedSelectionMergeMode.values()), value));
+ }
+ }
+
+ /// Returns the minimum sorted-segment fraction for [SortedSelectionMergeMode#AUTO], or null if not set.
+ @Nullable
+ public static Double getSortedSelectionMergeAutoMinSortedRatio(Map<String, String> queryOptions) {
+ return checkedParseRatio(QueryOptionKey.SORTED_SELECTION_MERGE_AUTO_MIN_SORTED_RATIO,
+ queryOptions.get(QueryOptionKey.SORTED_SELECTION_MERGE_AUTO_MIN_SORTED_RATIO));
}
@Nullable
@@ -902,6 +923,29 @@ public class QueryOptionsUtils {
return value;
}
+ /// Parses a fraction in the closed range [0, 1].
+ @Nullable
+ private static Double checkedParseRatio(String optionName, @Nullable String optionValue) {
+ if (optionValue == null) {
+ return null;
+ }
+ double value;
+ try {
+ value = Double.parseDouble(optionValue.trim());
+ } catch (NumberFormatException nfe) {
+ throw ratioParseException(optionName, optionValue);
+ }
+ if (!Double.isFinite(value) || value < 0 || value > 1) {
+ throw ratioParseException(optionName, optionValue);
+ }
+ return value;
+ }
+
+ private static IllegalArgumentException ratioParseException(String optionName, String optionValue) {
+ return new IllegalArgumentException(
+ String.format("%s must be a number between 0 and 1, got: %s", optionName, optionValue));
+ }
+
@Nullable
private static Long checkedParseLongPositive(String optionName, @Nullable String optionValue) {
return checkedParseLong(optionName, optionValue, 1);
diff --git a/pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java b/pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java
index ffb698e789d..ccd696f1c5c 100644
--- a/pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java
+++ b/pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java
@@ -25,6 +25,7 @@ import java.util.List;
import java.util.Map;
import java.util.Set;
import org.apache.pinot.spi.config.table.FieldConfig;
+import org.apache.pinot.spi.utils.CommonConstants.Server.SortedSelectionMergeMode;
import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
@@ -53,6 +54,54 @@ public class QueryOptionsUtilsTest {
private static final List<String> POSITIVE_LONG_KEYS =
List.of(TIMEOUT_MS, MAX_SERVER_RESPONSE_SIZE_BYTES, MAX_QUERY_RESPONSE_SIZE_BYTES);
+ @Test
+ public void shouldParseSortedSelectionMergeMode() {
+ assertNull(QueryOptionsUtils.getSortedSelectionMergeMode(Map.of()),
+ "An unset mode must stay null so the caller keeps its default");
+ assertEquals(QueryOptionsUtils.getSortedSelectionMergeMode(Map.of(SORTED_SELECTION_MERGE_MODE, "on")),
+ SortedSelectionMergeMode.ON);
+ assertEquals(QueryOptionsUtils.getSortedSelectionMergeMode(Map.of(SORTED_SELECTION_MERGE_MODE, "OFF")),
+ SortedSelectionMergeMode.OFF);
+ assertEquals(QueryOptionsUtils.getSortedSelectionMergeMode(Map.of(SORTED_SELECTION_MERGE_MODE, " Auto ")),
+ SortedSelectionMergeMode.AUTO, "The value must be trimmed and matched case-insensitively");
+ }
+
+ @Test
+ public void shouldRejectUnknownSortedSelectionMergeMode() {
+ try {
+ QueryOptionsUtils.getSortedSelectionMergeMode(Map.of(SORTED_SELECTION_MERGE_MODE, "yes"));
+ fail("Expected an unknown mode to be rejected rather than silently ignored");
+ } catch (IllegalArgumentException e) {
+ assertTrue(e.getMessage().contains(SORTED_SELECTION_MERGE_MODE) && e.getMessage().contains("[OFF, ON, AUTO]"),
+ "The message must name the option and list the valid values, got: " + e.getMessage());
+ }
+ }
+
+ @Test
+ public void shouldParseSortedSelectionMergeAutoMinSortedRatio() {
+ assertNull(QueryOptionsUtils.getSortedSelectionMergeAutoMinSortedRatio(Map.of()));
+ assertEquals(
+ QueryOptionsUtils.getSortedSelectionMergeAutoMinSortedRatio(
+ Map.of(SORTED_SELECTION_MERGE_AUTO_MIN_SORTED_RATIO, "0.75")), 0.75);
+ // The closed range is meaningful: 0 means "any segment set" and 1 means "every segment must be sorted".
+ assertEquals(
+ QueryOptionsUtils.getSortedSelectionMergeAutoMinSortedRatio(
+ Map.of(SORTED_SELECTION_MERGE_AUTO_MIN_SORTED_RATIO, "0")), 0.0);
+ assertEquals(
+ QueryOptionsUtils.getSortedSelectionMergeAutoMinSortedRatio(
+ Map.of(SORTED_SELECTION_MERGE_AUTO_MIN_SORTED_RATIO, "1")), 1.0);
+ for (String invalid : List.of("-0.1", "1.1", "abc", "NaN")) {
+ try {
+ QueryOptionsUtils.getSortedSelectionMergeAutoMinSortedRatio(
+ Map.of(SORTED_SELECTION_MERGE_AUTO_MIN_SORTED_RATIO, invalid));
+ fail("Expected the ratio '" + invalid + "' to be rejected");
+ } catch (IllegalArgumentException e) {
+ assertTrue(e.getMessage().contains(SORTED_SELECTION_MERGE_AUTO_MIN_SORTED_RATIO),
+ "The message must name the option, got: " + e.getMessage());
+ }
+ }
+ }
+
@Test
public void shouldConvertCaseInsensitiveMapToUseCorrectValues() {
// Given:
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperator.java
index b99d6f0edc7..527e7a77fb7 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperator.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperator.java
@@ -21,10 +21,8 @@ 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;
@@ -32,7 +30,6 @@ 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;
@@ -40,13 +37,9 @@ 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;
@@ -77,7 +70,8 @@ import org.slf4j.LoggerFactory;
/// (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.
+/// remaining segments are never acquired or read. See {@link #activateEligibleCursors(SegmentCursor)} 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.
///
@@ -92,15 +86,19 @@ import org.slf4j.LoggerFactory;
/// {@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.
+/// **Streaming paths only.** This operator is installed only where a
+/// {@link org.apache.pinot.core.query.executor.ResultsBlockStreamer} is present (the
+/// MSE leaf, driven by {@link org.apache.pinot.core.operator.streaming.StreamingInstanceResponseOperator}, and the
+/// gRPC streaming endpoint). The merge emits many bounded {@link SelectionResultsBlock}s from successive
+/// {@link #getNextBlock()} calls followed by a final {@link MetadataResultsBlock}. On the blocking single-stage
+/// path the caller expects one complete block from a single {@link #getNextBlock()} call, so nothing here could
+/// stream: the merge would accumulate the whole {@code limit + offset} result while holding every activated cursor's
+/// materialized block alive. That path keeps {@link MinMaxValueBasedSelectionOrderByCombineOperator}.
///
/// **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
+/// not use the base worker-queue model. The base entry points that it replaces - {@link #processSegments()} and
+/// {@link #isQuerySatisfied(SelectionResultsBlock, Object)} - are 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
@@ -110,7 +108,6 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
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;
@@ -124,7 +121,8 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
private int _numRowsEmitted;
private List<Object[]> _outputRows;
private boolean _done;
- /// Captured from the first child block seen; all child blocks share the same schema
+ /// 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
@@ -134,35 +132,49 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
private List<String> _unreportedDataSchemaMismatchErrors;
public StreamingSelectionOrderByCombineOperator(List<Operator> operators, QueryContext queryContext,
- ExecutorService executorService, boolean streaming) {
+ 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);
- _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();
+ _blockSize = queryContext.getSortedSelectionMergeBlockSize();
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());
+ // 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 && !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);
+ if (firstOrderByColumn == null) {
+ _sortedCursors[i] = new SegmentCursor(operator, null, null);
+ continue;
+ }
DataSourceMetadata metadata =
operator.getIndexSegment().getDataSource(firstOrderByColumn, queryContext.getSchema())
.getDataSourceMetadata();
_sortedCursors[i] = new SegmentCursor(operator, metadata.getMinValue(), metadata.getMaxValue());
}
- sortCursorsByMinMax();
+ if (firstOrderByColumn != null) {
+ sortCursorsByMinMax();
+ }
_priorityQueue = new PriorityQueue<>(Math.max(1, _numOperators),
(o1, o2) -> _comparator.compare(o1.currentHead(), o2.currentHead()));
@@ -222,6 +234,14 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
"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) {
@@ -230,41 +250,63 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
}
try {
long endTimeMs = _queryContext.getEndTimeMs();
+ // The leader is retained outside the heap: segments are usually near-disjoint on a time-like leading order-by
+ // column, so one cursor supplies long runs and re-heaping it per row costs two O(log k) sifts to reach the same
+ // cursor again. While it keeps winning we only peek and compare.
+ // TODO: move whole runs with one addAll(subList(...)) after a binary search for the cut. Deferred because
+ // advancing _numRowsEmitted by a run steps over the 0x1FFF mask in the termination check below, losing
+ // cancellation/pause/deadline observation; needs an unconditional check per run or a run-length clamp.
+ // See https://github.com/apache/pinot/pull/19120#discussion_r3871714056
+ SegmentCursor leader = null;
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();
- if (cursor == null) {
- // All active cursors exhausted (and pruning guarantees the rest cannot contribute).
- break;
+ activateEligibleCursors(leader);
+ // Choose the leader strictly AFTER activation, so a cursor just activated can win instead of being skipped
+ // for one row.
+ if (leader == null) {
+ leader = _priorityQueue.poll();
+ if (leader == null) {
+ // All active cursors exhausted (and pruning guarantees the rest cannot contribute).
+ break;
+ }
+ } else {
+ SegmentCursor top = _priorityQueue.peek();
+ if (top != null && _comparator.compare(leader.currentHead(), top.currentHead()) > 0) {
+ // Lost the lead; hand it back and take the new minimum. A tie keeps the incumbent, reordering only rows
+ // equal on every order-by expression -- the heap's own tie-breaking is arbitrary too.
+ _priorityQueue.offer(leader);
+ leader = _priorityQueue.poll();
+ }
}
- _outputRows.add(cursor.currentHead());
+ _outputRows.add(leader.currentHead());
_numRowsEmitted++;
- cursor.advance();
- if (cursor.currentHead() != null) {
- _priorityQueue.offer(cursor);
+ leader.advance();
+ if (leader.currentHead() == null) {
+ leader = null;
}
- if (_streaming && _outputRows.size() >= _blockSize) {
+ if (_outputRows.size() >= _blockSize) {
+ // Restore the invariant before handing control back: at every entry to and exit from getNextBlock() the
+ // heap holds all active cursors, so nothing else has to know about the retained leader.
+ if (leader != null) {
+ _priorityQueue.offer(leader);
+ }
return flushDataBlock();
}
}
- // Merge complete.
+ // Merge complete; the leader goes back so finish() sees a consistent heap.
+ if (leader != null) {
+ _priorityQueue.offer(leader);
+ }
finish();
if (!_outputRows.isEmpty()) {
- // Streaming: the final partial data block (next call returns the terminal metadata block).
- // Single-stage: the single complete block.
+ // The final partial data block; the next call returns the terminal metadata block.
return flushDataBlock();
}
- if (_streaming) {
- return attachDataSchemaMismatchErrors(attachExecutionStats(new MetadataResultsBlock()));
- }
- // Single-stage with no rows: still return a (single) block carrying the schema and execution stats.
- return attachDataSchemaMismatchErrors(attachExecutionStats(
- new SelectionResultsBlock(resolveDataSchema(), List.of(), _comparator, _queryContext)));
+ return attachDataSchemaMismatchErrors(attachExecutionStats(new MetadataResultsBlock()));
} catch (Exception e) {
_done = true;
releaseAllCursors();
@@ -278,32 +320,26 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
}
}
- /// Activates not-yet-active cursors whose min/max value can still contribute before the current merge frontier.
+ /// Activates not-yet-active cursors whose min/max bound can still reach the merge frontier.
+ ///
+ /// Cursors are visited in min/max order and {@code _nextToActivate} advances only on a real activation, so a
+ /// {@code break} defers the current cursor to a later call with a risen frontier rather than skipping it. Deferring
+ /// is safe because the first order-by column is the primary sort key: a cursor whose range starts past the frontier
+ /// holds no row sorting before it. With no frontier known yet, activation is forced.
///
- /// Cursors are visited in min/max-sorted order and {@code _nextToActivate} is advanced only when a cursor is
- /// actually activated; a {@code break} merely defers the current cursor, which is re-evaluated against the (rising)
- /// frontier on every subsequent call. Correctness: a not-yet-activated cursor whose first order-by value range starts
- /// strictly beyond the current heap head cannot contain a row that sorts before that head (the first order-by column
- /// is the primary sort key), so the head is the true global minimum and is safe to emit; the deferred cursor is
- /// activated later, exactly when the frontier reaches its min/max. When the heap is empty the frontier is unknown, so
- /// activation is forced (never pruned), which also drains any segments that sort entirely after the ones seen so far.
- private void activateEligibleCursors() {
+ /// The frontier is the row about to be emitted. Understating it defers a cursor that could have supplied that row
+ /// and the merge emits out of order, where overstating it only activates a segment early. The leader is retained
+ /// outside the heap, so that row is either its head or the heap head: a bound past both is past the frontier.
+ private void activateEligibleCursors(@Nullable SegmentCursor leader) {
while (_nextToActivate < _sortedCursors.length) {
SegmentCursor cursor = _sortedCursors[_nextToActivate];
if (_pruningEnabled) {
Comparable bound = _asc ? cursor._minValue : cursor._maxValue;
- // A null bound means the segment must always be processed. Otherwise, only prune against a non-null frontier;
- // if the head's first order-by value is null we cannot compare, so fall through and activate.
- if (bound != null && !_priorityQueue.isEmpty()) {
- Object headValue = _priorityQueue.peek().currentHead()[0];
- if (headValue != null) {
- // Both come from the same first order-by column: the metadata min/max and the materialized row[0] share
- // the column's stored type, so this comparison is type-safe (same assumption as MinMaxValueBased...).
- int cmp = bound.compareTo(headValue);
- if (_asc ? cmp > 0 : cmp < 0) {
- break;
- }
- }
+ // Re-read per iteration: each activation below can offer a smaller head into the heap.
+ SegmentCursor top = _priorityQueue.peek();
+ // A null bound always activates, and with no candidate at all there is no frontier to prune against.
+ if (bound != null && (leader != null || top != null) && sortsBeyond(bound, leader) && sortsBeyond(bound, top)) {
+ break;
}
}
cursor.activate();
@@ -314,6 +350,24 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
}
}
+ /// Whether `bound` sorts past the cursor's head on the first order-by column, so that segment cannot supply the
+ /// head. An absent cursor imposes no constraint; a null head cannot be compared, so this reports `false` and the
+ /// caller activates. Tests column 0 only, as the pruning bound always has, rather than the full-row comparator --
+ /// which would put every order-by column on the per-row path.
+ private boolean sortsBeyond(Comparable bound, @Nullable SegmentCursor cursor) {
+ if (cursor == null) {
+ return true;
+ }
+ Object headValue = cursor.currentHead()[0];
+ if (headValue == null) {
+ return false;
+ }
+ // Both come from the same first order-by column: the metadata min/max and the materialized row[0] share the
+ // column's stored type, so this comparison is type-safe (same assumption as MinMaxValueBased...).
+ int cmp = bound.compareTo(headValue);
+ return _asc ? cmp > 0 : cmp < 0;
+ }
+
/// Marks the merge done and releases any segments still held by un-drained cursors (e.g. when LIMIT is reached).
private void finish() {
_done = true;
@@ -327,14 +381,16 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
}
/// Returns the accumulated output rows as a sorted {@link SelectionResultsBlock} and resets the output buffer. The
- /// block carries the comparator so the broker-side n-way reduce stays correct. In single-stage mode it is the only
- /// block, so execution stats are attached; in streaming mode stats are attached to the terminal metadata block.
+ /// block carries the comparator so the broker-side n-way reduce stays correct. Execution stats are not attached
+ /// here: they go on the terminal metadata block.
private BaseResultsBlock flushDataBlock() {
List<Object[]> rows = _outputRows;
_outputRows = newOutputList();
- SelectionResultsBlock block = new SelectionResultsBlock(resolveDataSchema(), rows, _comparator, _queryContext);
+ // Rows only reach the output buffer through pullBlock(), which captures the schema before accepting any of them.
+ assert _dataSchema != null;
+ SelectionResultsBlock block = new SelectionResultsBlock(_dataSchema, rows, _comparator, _queryContext);
attachDataSchemaMismatchErrors(block);
- return _streaming ? block : attachExecutionStats(block);
+ return block;
}
/// Records that a segment's block was dropped because its schema disagreed with the merge schema. Deduplicated by
@@ -372,37 +428,6 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
return new ArrayList<>(Math.max(1, capacity));
}
- /// Returns the data schema captured from the first child block. If no segment produced a block (every segment is a
- /// streaming operator that matched zero rows), reconstructs the schema from the first segment so that an empty result
- /// still carries a valid, correctly-ordered schema (order-by expressions first, matching the child blocks' layout).
- private DataSchema resolveDataSchema() {
- if (_dataSchema != null) {
- return _dataSchema;
- }
- IndexSegment indexSegment = _operators.get(0).getIndexSegment();
- List<ExpressionContext> expressions = SelectionOperatorUtils.extractExpressions(_queryContext, indexSegment);
- Set<String> columns = new HashSet<>();
- for (ExpressionContext expression : expressions) {
- expression.getColumns(columns);
- }
- Map<String, DataSource> dataSourceMap = new HashMap<>();
- for (String column : columns) {
- dataSourceMap.put(column, indexSegment.getDataSource(column, _queryContext.getSchema()));
- }
- int numExpressions = expressions.size();
- String[] columnNames = new String[numExpressions];
- ColumnDataType[] columnDataTypes = new ColumnDataType[numExpressions];
- for (int i = 0; i < numExpressions; i++) {
- ExpressionContext expression = expressions.get(i);
- columnNames[i] = expression.toString();
- TransformFunction transformFunction = TransformFunctionFactory.get(expression, dataSourceMap);
- columnDataTypes[i] = ColumnDataType.fromDataType(transformFunction.getResultMetadata().getDataType(),
- transformFunction.getResultMetadata().isSingleValue());
- }
- _dataSchema = new DataSchema(columnNames, columnDataTypes);
- return _dataSchema;
- }
-
/// Iterates a single segment's locally-sorted rows, pulling blocks lazily from its operator. Streaming-backed cursors
/// loop until the operator returns {@code null}; single-block-backed cursors read exactly one block. The segment is
/// acquired on activation and released once exhausted or when the combine finishes (see
@@ -415,6 +440,7 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
private final Comparable _maxValue;
private boolean _streamingChild;
+ private boolean _streamingChildResolved;
private List<Object[]> _rows;
private int _pos;
private boolean _acquired;
@@ -439,7 +465,6 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
/// nothing.
void activate() {
acquireSegment();
- _streamingChild = isStreamingChild();
if (!pullBlock()) {
exhaust();
}
@@ -488,8 +513,9 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
_head = rows.get(0);
return true;
}
- // Defensive: an unexpected empty (non-null) block. Keep pulling only for streaming children; a single-block
- // child yields exactly one block, so treat it as exhausted.
+ // An empty (non-null) block: a streaming child emits one on a zero-match segment to deliver its schema, which
+ // the capture above has just taken. Keep pulling only for streaming children; a single-block child yields
+ // exactly one block, so treat it as exhausted.
if (!_streamingChild) {
return false;
}
@@ -498,21 +524,30 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
private SelectionResultsBlock nextBlock() {
try {
- return (SelectionResultsBlock) _operator.nextBlock();
+ SelectionResultsBlock block = (SelectionResultsBlock) _operator.nextBlock();
+ if (!_streamingChildResolved) {
+ // The call above materializes the wrapped child, so its type can be read without forcing a separate
+ // materializeChildOperator(). Resolved after the first pull rather than before it because that is the
+ // earliest point at which the child exists.
+ _streamingChild = isStreamingChild();
+ _streamingChildResolved = true;
+ }
+ return block;
} catch (RuntimeException e) {
throw wrapOperatorException(_operator, e);
}
}
/// Returns whether the underlying child operator is the lazy {@link StreamingSelectionOrderByOperator}. Must be
- /// called after {@link #acquireSegment()} because materializing the wrapped child runs the plan node, which
- /// accesses segment buffers.
+ /// called after the first {@link Operator#nextBlock()}, which is what materializes a wrapped child.
private boolean isStreamingChild() {
Operator underlying = _operator;
if (_operator instanceof AcquireReleaseColumnsSegmentOperator) {
- AcquireReleaseColumnsSegmentOperator wrapper = (AcquireReleaseColumnsSegmentOperator) _operator;
- wrapper.materializeChildOperator();
- underlying = wrapper.getChildOperators().get(0);
+ // The child is materialized by AcquireReleaseColumnsSegmentOperator#getNextBlock() before it delegates, so it
+ // is non-null by the time we get here. Left unchecked deliberately: getChildOperators() returns
+ // List.of(child), which itself throws NullPointerException on a null element, so a refactor that breaks the
+ // invariant already fails there and an explicit check here would be unreachable.
+ underlying = ((AcquireReleaseColumnsSegmentOperator) _operator).getChildOperators().get(0);
}
return underlying instanceof StreamingSelectionOrderByOperator;
}
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperator.java
index e66a6ee0391..10dbd97ccde 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperator.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperator.java
@@ -20,7 +20,6 @@ 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;
@@ -29,11 +28,13 @@ 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;
@@ -48,6 +49,7 @@ 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;
@@ -65,6 +67,9 @@ import org.roaringbitmap.RoaringBitmap;
/// 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))}).
///
+/// Every instance emits at least one block before {@code null}: a segment matching no rows emits a single empty block
+/// carrying the {@link 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** ({@code numSortedExpressions == numOrderByExpressions}, e.g. {@code ORDER BY sorted}):
@@ -118,6 +123,8 @@ public class StreamingSelectionOrderByOperator extends BaseOperator<SelectionRes
/// 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;
@@ -131,6 +138,8 @@ public class StreamingSelectionOrderByOperator extends BaseOperator<SelectionRes
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;
@@ -200,8 +209,19 @@ public class StreamingSelectionOrderByOperator extends BaseOperator<SelectionRes
List<Object[]> rows = _tailToSort ? nextRun() : nextSortedRows();
if (rows == null || rows.isEmpty()) {
_exhausted = true;
- return null;
+ 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);
}
@@ -211,14 +231,16 @@ public class StreamingSelectionOrderByOperator extends BaseOperator<SelectionRes
}
/// 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.
+ /// trimmed to the remaining {@code limit + offset} budget. Returns {@code 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 = _projectOperator.nextBlock();
+ ValueBlock valueBlock = nextNonEmptyBlock();
if (valueBlock == null) {
return null;
}
@@ -297,40 +319,54 @@ public class StreamingSelectionOrderByOperator extends BaseOperator<SelectionRes
/// is exhausted.
@Nullable
private Object[] nextRow() {
- while (true) {
- if (_currentBlock == null || _currentPos >= _currentNumDocs) {
- if (_projectExhausted) {
- return null;
- }
- _currentBlock = _projectOperator.nextBlock();
- if (_currentBlock == null) {
- _projectExhausted = true;
- return null;
- }
- BlockValSet[] blockValSets = new BlockValSet[_numPhase1Columns];
+ 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++) {
- blockValSets[i] = _currentBlock.getBlockValueSet(_phase1Expressions.get(i));
- }
- _currentFetcher = new RowBasedBlockValueFetcher(blockValSets);
- _currentNumDocs = _currentBlock.getNumDocs();
- _currentDocIds = _twoPhase ? _currentBlock.getDocIds() : null;
- if (_nullHandlingEnabled) {
- _currentNullBitmaps = new RoaringBitmap[_numPhase1Columns];
- for (int i = 0; i < _numPhase1Columns; i++) {
- _currentNullBitmaps[i] = blockValSets[i].getNullBitmap();
- }
- }
- _currentPos = 0;
- _numDocsScanned += _currentNumDocs;
- _numEntriesScannedPostFilter += (long) _currentNumDocs * _projectOperator.getNumColumnsProjected();
- reportScanCost(_currentNumDocs, (long) _currentNumDocs * _projectOperator.getNumColumnsProjected());
- if (_currentNumDocs == 0) {
- _currentBlock = null;
- continue;
+ _currentNullBitmaps[i] = blockValSets[i].getNullBitmap();
}
}
- int rowId = _currentPos++;
- return materializeRow(_currentFetcher, _currentDocIds, _currentNullBitmaps, rowId);
+ _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 {@code 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;
+ }
}
}
@@ -352,15 +388,23 @@ public class StreamingSelectionOrderByOperator extends BaseOperator<SelectionRes
return row;
}
- /// Drains a max-heap (created with the reversed comparator) into an ascending list, mutable so the second pass can
- /// fill non-order-by values in place.
+ /// 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();
- Object[][] sortedRows = new Object[numRows][];
- for (int i = numRows - 1; i >= 0; i--) {
- sortedRows[i] = heap.poll();
+ // 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());
}
- return Arrays.asList(sortedRows);
+ 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
@@ -417,7 +461,7 @@ public class StreamingSelectionOrderByOperator extends BaseOperator<SelectionRes
}
if (_dataSchema == null) {
- _dataSchema = buildTwoPhaseDataSchema(transformOperator);
+ _dataSchema = buildTwoPhaseDataSchema(transformOperator::getResultColumnContext);
}
}
}
@@ -433,7 +477,24 @@ public class StreamingSelectionOrderByOperator extends BaseOperator<SelectionRes
return new DataSchema(columnNames, columnDataTypes);
}
- private DataSchema buildTwoPhaseDataSchema(TransformOperator transformOperator) {
+ /// Resolves a non-order-by expression's result type without building the phase-2 pipeline, for the zero-match case
+ /// where there is nothing to fetch. This reproduces {@link TransformOperator#getResultColumnContext} exactly -- that
+ /// method resolves against its project operator's source column contexts, which for phase 2 are
+ /// {@link #_phase2DataSourceMap} -- so the schema is identical to the one the fetch path builds, at metadata cost
+ /// only. Worth the indirection because a selective filter can leave many segments of a table matching nothing, and
+ /// each would otherwise stand up a full projection and transform operator to read types it already knows.
+ private ColumnContext resolveResultColumnContext(ExpressionContext expression) {
+ if (_phase2ColumnContextMap == null) {
+ _phase2ColumnContextMap = new HashMap<>(HashUtil.getHashMapCapacity(_phase2DataSourceMap.size()));
+ for (Map.Entry<String, DataSource> entry : _phase2DataSourceMap.entrySet()) {
+ _phase2ColumnContextMap.put(entry.getKey(), ColumnContext.fromDataSource(entry.getValue()));
+ }
+ }
+ return ColumnContext.fromTransformFunction(
+ TransformFunctionFactory.get(expression, _phase2ColumnContextMap, _queryContext));
+ }
+
+ private DataSchema buildTwoPhaseDataSchema(Function<ExpressionContext, ColumnContext> resultColumnContexts) {
int numNonOrderByExpressions = _nonOrderByExpressions.size();
String[] columnNames = new String[_numExpressions];
DataSchema.ColumnDataType[] columnDataTypes = new DataSchema.ColumnDataType[_numExpressions];
@@ -445,7 +506,7 @@ public class StreamingSelectionOrderByOperator extends BaseOperator<SelectionRes
_orderByColumnContexts[i].isSingleValue());
}
for (int i = 0; i < numNonOrderByExpressions; i++) {
- ColumnContext columnContext = transformOperator.getResultColumnContext(_nonOrderByExpressions.get(i));
+ ColumnContext columnContext = resultColumnContexts.apply(_nonOrderByExpressions.get(i));
columnDataTypes[_numOrderByExpressions + i] =
DataSchema.ColumnDataType.fromDataType(columnContext.getDataType(), columnContext.isSingleValue());
}
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/plan/CombinePlanNode.java b/pinot-core/src/main/java/org/apache/pinot/core/plan/CombinePlanNode.java
index 298897e3ec2..76afc21d938 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/plan/CombinePlanNode.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/plan/CombinePlanNode.java
@@ -132,16 +132,16 @@ public class CombinePlanNode implements PlanNode {
// Use streaming operator only for non-empty selection-only query
return createStreamingSelectionOnlyCombineOperator(operators);
}
- // 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.
+ // Streaming selection order-by (opt-in via the sortedSelectionMergeMode hint). Selection-only already returned
+ // above, so reaching here with a non-empty limit and an order-by present implies selection order-by.
+ //
+ // No sortedness check here: InstancePlanMakerImplV2 has already resolved AUTO against segment metadata, and a
+ // leading order-by expression is deliberately allowed because that is what ON is for. The merge handles a
+ // non-identifier leading expression by giving up frontier pruning, see
+ // StreamingSelectionOrderByCombineOperator.
if (_queryContext.isSortedSelectionMergeEnabled() && QueryContextUtils.isSelectionQuery(_queryContext)
- && _queryContext.getLimit() != 0) {
- List<OrderByExpressionContext> orderByExpressions = _queryContext.getOrderByExpressions();
- if (orderByExpressions != null
- && orderByExpressions.get(0).getExpression().getType() == ExpressionContext.Type.IDENTIFIER) {
- return new StreamingSelectionOrderByCombineOperator(operators, _queryContext, _executorService,
- true /* streaming */);
- }
+ && _queryContext.getLimit() != 0 && _queryContext.getOrderByExpressions() != null) {
+ return new StreamingSelectionOrderByCombineOperator(operators, _queryContext, _executorService);
}
// Streaming flushes partial results, so it needs an aggregation above to merge them back together.
// Leaves that must return final results are excluded, see StreamingGroupByCombineOperator.
@@ -198,10 +198,6 @@ public class CombinePlanNode implements PlanNode {
List<OrderByExpressionContext> orderByExpressions = _queryContext.getOrderByExpressions();
assert orderByExpressions != null;
if (orderByExpressions.get(0).getExpression().getType() == ExpressionContext.Type.IDENTIFIER) {
- if (_queryContext.isSortedSelectionMergeEnabled()) {
- return new StreamingSelectionOrderByCombineOperator(operators, _queryContext, _executorService,
- false /* streaming */);
- }
return new MinMaxValueBasedSelectionOrderByCombineOperator(operators, _queryContext, _executorService);
}
return new SelectionOrderByCombineOperator(operators, _queryContext, _executorService);
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/plan/SelectionPlanNode.java b/pinot-core/src/main/java/org/apache/pinot/core/plan/SelectionPlanNode.java
index 4849a193eeb..f2a39b14777 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/plan/SelectionPlanNode.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/plan/SelectionPlanNode.java
@@ -77,7 +77,7 @@ public class SelectionPlanNode implements PlanNode {
int numOrderByExpressions = orderByExpressions.size();
// Although it is a break of abstraction, some code, specially merging, assumes that if there is an order by
// expression the operator will return a block whose selection result is a priority queue.
- int sortedColumnsPrefixSize = getSortedColumnsPrefix(orderByExpressions, _queryContext.isNullHandlingEnabled());
+ int sortedColumnsPrefixSize = getSortedColumnsPrefix(orderByExpressions);
if (sortedColumnsPrefixSize > 0) {
int maxDocsPerCall = DocIdSetPlanNode.MAX_DOC_PER_CALL;
// The first order by expressions are sorted (either asc or desc).
@@ -95,15 +95,23 @@ public class SelectionPlanNode implements PlanNode {
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.
+ // lazily. Reaching here means the leading order-by is already a sorted prefix, so it is either a LITERAL or a
+ // physically sorted IDENTIFIER; only the identifier case has a forward index to scan in order. The combine-side
+ // gate deliberately does not repeat this check -- it accepts materialized children too, which is how
+ // SortedSelectionMergeMode.ON forces the path for an expression order-by.
+ //
+ // The DESC-incompatible sorted case still falls back to the materialized SelectionPartiallyOrderedByDescOperation
+ // below so global order stays correct.
+
+ // Set when the streaming attempt below builds a project the materialized fallback can take over.
+ BaseProjectOperator<?> reusableSortedByProject = null;
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.
+ boolean projectsAllExpressions = expressions.size() <= numOrderByExpressions;
List<ExpressionContext> projectExpressions = expressions;
- if (expressions.size() > numOrderByExpressions) {
+ if (!projectsAllExpressions) {
projectExpressions = new ArrayList<>(numOrderByExpressions);
for (OrderByExpressionContext orderByExpression : orderByExpressions) {
projectExpressions.add(orderByExpression.getExpression());
@@ -115,10 +123,18 @@ public class SelectionPlanNode implements PlanNode {
return new StreamingSelectionOrderByOperator(_indexSegment, _queryContext, expressions,
streamingProjectOperator, sortedColumnsPrefixSize);
}
- // DESC-incompatible: fall through to the materialized fallback (rebuilds the project over all expressions).
+ // The project cannot scan in the query's direction, which only a DESC query can hit. Fall through to the
+ // materialized fallback, reusing this project when it already covers the full expression list: that is what
+ // the fallback would rebuild, and it was never driven. A narrower one is dropped unclosed, per the standing
+ // TODO on ProjectPlanNode#run.
+ if (projectsAllExpressions) {
+ reusableSortedByProject = streamingProjectOperator;
+ }
}
- BaseProjectOperator<?> projectOperator = getSortedByProject(expressions, maxDocsPerCall, orderByExpressions);
+ BaseProjectOperator<?> projectOperator = reusableSortedByProject != null
+ ? reusableSortedByProject
+ : getSortedByProject(expressions, maxDocsPerCall, orderByExpressions);
if (projectOperator.isCompatibleWith(queryOrder)) {
return new SelectionPartiallyOrderedByLinearOperator(_indexSegment, _queryContext, expressions, projectOperator,
sortedColumnsPrefixSize);
@@ -181,10 +197,10 @@ public class SelectionPlanNode implements PlanNode {
///
/// @return the max number that guarantees that from the first expression to the returned number, the index is already
/// sorted.
- private int getSortedColumnsPrefix(List<OrderByExpressionContext> orderByExpressions, boolean isNullHandlingEnabled) {
+ private int getSortedColumnsPrefix(List<OrderByExpressionContext> orderByExpressions) {
boolean asc = orderByExpressions.get(0).isAsc();
for (int i = 0; i < orderByExpressions.size(); i++) {
- if (!isSorted(orderByExpressions.get(i), asc, isNullHandlingEnabled)) {
+ if (!isSorted(orderByExpressions.get(i), asc)) {
return i;
}
}
@@ -192,7 +208,7 @@ public class SelectionPlanNode implements PlanNode {
return orderByExpressions.size();
}
- private boolean isSorted(OrderByExpressionContext orderByExpression, boolean asc, boolean isNullHandlingEnabled) {
+ private boolean isSorted(OrderByExpressionContext orderByExpression, boolean asc) {
switch (orderByExpression.getExpression().getType()) {
case LITERAL: {
return true;
@@ -201,16 +217,7 @@ public class SelectionPlanNode implements PlanNode {
if (!orderByExpression.isAsc() == asc) {
return false;
}
- String column = orderByExpression.getExpression().getIdentifier();
- DataSource dataSource = _indexSegment.getDataSource(column, _queryContext.getSchema());
- // If there are null values, we cannot trust DataSourceMetadata.isSorted
- if (isNullHandlingEnabled) {
- NullValueVectorReader nullValueVector = dataSource.getNullValueVector();
- if (nullValueVector != null && !nullValueVector.getNullBitmap().isEmpty()) {
- return false;
- }
- }
- return dataSource.getDataSourceMetadata().isSorted();
+ return isColumnSorted(_indexSegment, _queryContext, orderByExpression.getExpression().getIdentifier());
}
case FUNCTION: // we could optimize monotonically increasing functions
default: {
@@ -218,4 +225,32 @@ public class SelectionPlanNode implements PlanNode {
}
}
}
+
+ /// Returns whether `column` can be relied on to be sorted (ascending) for this query in `segment`.
+ ///
+ /// This is [#isColumnPhysicallySorted] plus the null caveat: once null handling is on, the physical order is not
+ /// the order the query asks for, so a column carrying nulls is not usable as sorted. Reading the null bitmap
+ /// touches the segment's mapped buffer, so this must only be called once the segment has been acquired -- which is
+ /// why [AcquireReleaseColumnsSegmentPlanNode] defers the whole plan build until after `acquire()`.
+ public static boolean isColumnSorted(IndexSegment segment, QueryContext queryContext, String column) {
+ DataSource dataSource = segment.getDataSource(column, queryContext.getSchema());
+ // If there are null values, we cannot trust DataSourceMetadata.isSorted
+ if (queryContext.isNullHandlingEnabled()) {
+ NullValueVectorReader nullValueVector = dataSource.getNullValueVector();
+ if (nullValueVector != null && !nullValueVector.getNullBitmap().isEmpty()) {
+ return false;
+ }
+ }
+ return isColumnPhysicallySorted(segment, queryContext, column);
+ }
+
+ /// Returns whether `column` is physically sorted (ascending) in `segment`, ignoring nulls.
+ ///
+ /// Reads [org.apache.pinot.segment.spi.datasource.DataSourceMetadata] only, so it touches no column buffer and
+ /// needs no segment acquire. That is what lets
+ /// [org.apache.pinot.core.plan.maker.InstancePlanMakerImplV2] resolve
+ /// [org.apache.pinot.spi.utils.CommonConstants.Server.SortedSelectionMergeMode#AUTO] before any plan node is built.
+ public static boolean isColumnPhysicallySorted(IndexSegment segment, QueryContext queryContext, String column) {
+ return segment.getDataSource(column, queryContext.getSchema()).getDataSourceMetadata().isSorted();
+ }
}
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java b/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java
index 25ff2a4db17..e0c8c541184 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java
@@ -55,7 +55,9 @@ import org.apache.pinot.segment.spi.FetchContext;
import org.apache.pinot.segment.spi.IndexSegment;
import org.apache.pinot.segment.spi.SegmentContext;
import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey;
import org.apache.pinot.spi.utils.CommonConstants.Server;
+import org.apache.pinot.spi.utils.CommonConstants.Server.SortedSelectionMergeMode;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -205,6 +207,16 @@ public class InstancePlanMakerImplV2 implements PlanMaker {
public Plan makeInstancePlan(List<SegmentContext> segmentContexts, QueryContext queryContext,
ExecutorService executorService) {
applyQueryOptions(queryContext);
+ // No ResultsBlockStreamer here, so nothing can stream. SelectionPlanNode gates the streaming leaf on this same
+ // option and cannot see the streamer, so clear it for the whole plan rather than let it be half-honored.
+ // Read the mode directly rather than isSortedSelectionMergeEnabled(): this (non-streaming) plan never resolves
+ // AUTO at all, and that accessor would warn and fall back to OFF on an unresolved AUTO. Reading the mode
+ // directly avoids the spurious warning and is precise about intent here.
+ if (queryContext.getSortedSelectionMergeMode() != SortedSelectionMergeMode.OFF) {
+ LOGGER.debug("Ignoring {}={} for a non-streaming query: the merge requires a streaming-capable path",
+ QueryOptionKey.SORTED_SELECTION_MERGE_MODE, queryContext.getSortedSelectionMergeMode());
+ queryContext.setSortedSelectionMergeMode(SortedSelectionMergeMode.OFF);
+ }
int numSegments = segmentContexts.size();
List<PlanNode> planNodes = new ArrayList<>(numSegments);
@@ -273,7 +285,15 @@ public class InstancePlanMakerImplV2 implements PlanMaker {
// Set streaming selection order-by options (opt-in; gated on selection queries to prevent accidental routing
// if a downstream guard is ever missed)
if (QueryContextUtils.isSelectionQuery(queryContext)) {
- queryContext.setSortedSelectionMergeEnabled(QueryOptionsUtils.isSortedSelectionMergeEnabled(queryOptions));
+ SortedSelectionMergeMode sortedSelectionMergeMode = QueryOptionsUtils.getSortedSelectionMergeMode(queryOptions);
+ if (sortedSelectionMergeMode != null) {
+ queryContext.setSortedSelectionMergeMode(sortedSelectionMergeMode);
+ }
+ Double sortedSelectionMergeAutoMinSortedRatio =
+ QueryOptionsUtils.getSortedSelectionMergeAutoMinSortedRatio(queryOptions);
+ if (sortedSelectionMergeAutoMinSortedRatio != null) {
+ queryContext.setSortedSelectionMergeAutoMinSortedRatio(sortedSelectionMergeAutoMinSortedRatio);
+ }
Integer sortedSelectionMergeBlockSize =
QueryOptionsUtils.getSortedSelectionMergeBlockSize(queryOptions);
if (sortedSelectionMergeBlockSize != null) {
@@ -386,6 +406,7 @@ public class InstancePlanMakerImplV2 implements PlanMaker {
public Plan makeStreamingInstancePlan(List<SegmentContext> segmentContexts, QueryContext queryContext,
ExecutorService executorService, ResultsBlockStreamer streamer) {
applyQueryOptions(queryContext);
+ resolveSortedSelectionMergeMode(segmentContexts, queryContext);
int numSegments = segmentContexts.size();
List<PlanNode> planNodes = new ArrayList<>(numSegments);
@@ -412,6 +433,94 @@ public class InstancePlanMakerImplV2 implements PlanMaker {
new StreamingInstanceResponsePlanNode(combinePlanNode, segmentContexts, fetchContexts, queryContext, streamer));
}
+ /// Resolves [SortedSelectionMergeMode#AUTO] to `ON` or `OFF` from segment metadata alone, before any plan node is
+ /// built. `ON` and `OFF` pass through untouched.
+ ///
+ /// This has to happen here rather than in the plan-node gates. [CombinePlanNode] builds every leaf operator before
+ /// it evaluates its own gate, so a decision taken there would arrive too late to stop streaming leaves from being
+ /// built underneath a non-streaming combine -- the half-honored pairing that
+ /// [org.apache.pinot.core.operator.combine.StreamingSelectionOrderByCombineOperator] exists to avoid. Deciding once
+ /// here leaves both gates reading a single settled value.
+ ///
+ /// The heuristic is the fraction of queried segments physically sorted on the leading ORDER BY column, compared
+ /// against [QueryContext#getSortedSelectionMergeAutoMinSortedRatio()]. A DESC leading expression additionally
+ /// requires `allowReverseOrder`, without which no leaf can stream in the query's direction. An unsorted child
+ /// forces a full scan plus top-K synchronously on the consumer thread, so it is the unsorted segments that cost
+ /// the query its `maxExecutionThreads` worth of parallelism. Reading
+ /// [org.apache.pinot.segment.spi.datasource.DataSourceMetadata] does not touch column buffers, so no segment is
+ /// acquired.
+ private void resolveSortedSelectionMergeMode(List<SegmentContext> segmentContexts, QueryContext queryContext) {
+ if (queryContext.getSortedSelectionMergeMode() != SortedSelectionMergeMode.AUTO) {
+ return;
+ }
+ queryContext.setSortedSelectionMergeMode(
+ isSortedEnoughForStreamingMerge(segmentContexts, queryContext) ? SortedSelectionMergeMode.ON
+ : SortedSelectionMergeMode.OFF);
+ }
+
+ private boolean isSortedEnoughForStreamingMerge(List<SegmentContext> segmentContexts, QueryContext queryContext) {
+ int numSegments = segmentContexts.size();
+ if (numSegments == 0) {
+ return false;
+ }
+ List<OrderByExpressionContext> orderByExpressions = queryContext.getOrderByExpressions();
+ if (orderByExpressions == null || orderByExpressions.isEmpty()) {
+ return false;
+ }
+ ExpressionContext firstOrderByExpression = orderByExpressions.get(0).getExpression();
+ if (firstOrderByExpression.getType() != ExpressionContext.Type.IDENTIFIER) {
+ // Sortedness is only defined for a physical column. Forcing the merge over an expression is what ON is for.
+ LOGGER.debug("Not using {} for a non-identifier leading ORDER BY expression: {}",
+ QueryOptionKey.SORTED_SELECTION_MERGE_MODE, firstOrderByExpression);
+ return false;
+ }
+ if (!orderByExpressions.get(0).isAsc() && !QueryOptionsUtils.isReverseOrderAllowed(
+ queryContext.getQueryOptions())) {
+ // DataSourceMetadata.isSorted() reports ascending physical order, so a DESC query can only stream when the
+ // project operator's docId scan is reversed, and SelectionPlanNode.getSortedByProject() only attempts that
+ // under allowReverseOrder. Without it every leaf falls back to the materialized
+ // SelectionPartiallyOrderedByDescOperation, and resolving to ON here would install the streaming combine over
+ // children that each materialize their top-K serially on the consumer thread -- losing both the MinMax
+ // combine's parallelism and its min/max segment pruning, in exchange for nothing.
+ //
+ // allowReverseOrder stays the user's lever rather than being implied here: ReverseDocIdSetOperator buffers the
+ // whole matching docId set into a RoaringBitmap before emitting its first block, which would trade away the
+ // bounded memory this merge exists to provide.
+ //
+ // TODO: allowReverseOrder=true still does not guarantee streaming leaves. getSortedByProject() swallows the
+ // failure when a docId set cannot be reversed and returns the unreversed operator, so the leaf gate then
+ // declines and AUTO has again resolved ON over materialized children. Same residual shape as the null
+ // handling case below; a shared "can this segment actually stream in the query's direction" predicate
+ // would close both. See https://github.com/apache/pinot/pull/19120#discussion_r3871713989
+ LOGGER.debug("Not using {} for a DESC leading ORDER BY expression without {}: {}",
+ QueryOptionKey.SORTED_SELECTION_MERGE_MODE, QueryOptionKey.ALLOW_REVERSE_ORDER, firstOrderByExpression);
+ return false;
+ }
+ String column = firstOrderByExpression.getIdentifier();
+ int numSorted = 0;
+ for (SegmentContext segmentContext : segmentContexts) {
+ // TODO: This deliberately uses the physical-sortedness predicate rather than the full
+ // SelectionPlanNode.isColumnSorted(), which additionally rejects a column carrying nulls when null
+ // handling is on. That null check reads the segment's mapped buffer via the null value vector, and this
+ // runs before any segment is acquired -- AcquireReleaseColumnsSegmentPlanNode defers the whole plan build
+ // precisely so that no planner touches a buffer pre-acquire. Consequence: with null handling on and a
+ // null-bearing leading column, AUTO can resolve to ON while the leaf gate then refuses to stream, giving
+ // a k-way merge over fully materialized children -- correct, but slower than the MinMax combine. Handle
+ // properly (an acquire-free nullability signal, or resolving AUTO per segment at acquire time) as a
+ // follow-up. See https://github.com/apache/pinot/pull/19120#discussion_r3871713975
+ if (SelectionPlanNode.isColumnPhysicallySorted(segmentContext.getIndexSegment(), queryContext, column)) {
+ numSorted++;
+ }
+ }
+ double sortedRatio = (double) numSorted / numSegments;
+ double minSortedRatio = queryContext.getSortedSelectionMergeAutoMinSortedRatio();
+ boolean enabled = sortedRatio >= minSortedRatio;
+ LOGGER.debug("{}=AUTO resolved to {}: {}/{} segments sorted on '{}' (ratio {}, threshold {})",
+ QueryOptionKey.SORTED_SELECTION_MERGE_MODE, enabled ? "ON" : "OFF", numSorted, numSegments, column,
+ sortedRatio, minSortedRatio);
+ return enabled;
+ }
+
@Override
public PlanNode makeStreamingSegmentPlanNode(SegmentContext segmentContext, QueryContext queryContext) {
if (QueryContextUtils.isSelectionOnlyQuery(queryContext) && queryContext.getLimit() != 0) {
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java b/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java
index 0d7f17775ce..384f800c8b2 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java
@@ -46,6 +46,9 @@ import org.apache.pinot.spi.config.table.FieldConfig;
import org.apache.pinot.spi.data.Schema;
import org.apache.pinot.spi.utils.CommonConstants.Broker;
import org.apache.pinot.spi.utils.CommonConstants.Server;
+import org.apache.pinot.spi.utils.CommonConstants.Server.SortedSelectionMergeMode;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
/// The `QueryContext` class encapsulates all the query related information extracted from the wiring Object.
@@ -66,6 +69,8 @@ import org.apache.pinot.spi.utils.CommonConstants.Server;
/// reduce the repetitive work for each segment.
@SuppressWarnings({"rawtypes", "unchecked"})
public class QueryContext {
+ private static final Logger LOGGER = LoggerFactory.getLogger(QueryContext.class);
+
private final String _tableName;
private final QueryContext _subquery;
private final List<ExpressionContext> _selectExpressions;
@@ -136,10 +141,17 @@ public class QueryContext {
private int _streamingGroupByFlushThreshold;
// Flush threshold for streaming distinct (0 = disabled)
private int _streamingDistinctFlushThreshold;
- /// Opt-in: use the streaming k-way-merge selection ORDER BY combine over sorted segments
- private boolean _sortedSelectionMergeEnabled;
+ /// Selects the streaming k-way-merge selection ORDER BY combine over sorted segments. `AUTO` is resolved to `ON`
+ /// or `OFF` by the plan maker before any plan node is built, so the plan-node gates only ever see a decided value.
+ private SortedSelectionMergeMode _sortedSelectionMergeMode = Server.DEFAULT_SORTED_SELECTION_MERGE_MODE;
+ /// Minimum fraction of sorted segments for `AUTO` to resolve to `ON`
+ private double _sortedSelectionMergeAutoMinSortedRatio =
+ Server.DEFAULT_SORTED_SELECTION_MERGE_AUTO_MIN_SORTED_RATIO;
/// Output block size (rows) for the streaming selection ORDER BY combine
private int _sortedSelectionMergeBlockSize = Broker.DEFAULT_SORTED_SELECTION_MERGE_BLOCK_SIZE;
+ // Guards the one-time warning in isSortedSelectionMergeEnabled() so that an unresolved AUTO does not log once per
+ // segment/combine call for the same query.
+ private volatile boolean _unresolvedSortedSelectionMergeModeWarned;
// Whether null handling is enabled
private boolean _nullHandlingEnabled;
@@ -515,12 +527,44 @@ public class QueryContext {
_streamingDistinctFlushThreshold = streamingDistinctFlushThreshold;
}
+ public SortedSelectionMergeMode getSortedSelectionMergeMode() {
+ return _sortedSelectionMergeMode;
+ }
+
+ public void setSortedSelectionMergeMode(SortedSelectionMergeMode sortedSelectionMergeMode) {
+ _sortedSelectionMergeMode = sortedSelectionMergeMode;
+ }
+
+ /// Whether the plan should use the streaming selection ORDER BY merge.
+ ///
+ /// `AUTO` is a request for a decision, not a decision, and collapsing it to "enabled" would force the merge
+ /// unconditionally while silently ignoring [#getSortedSelectionMergeAutoMinSortedRatio()] -- the half-honored
+ /// failure the single-resolution design exists to prevent. The plan maker resolves it against segment metadata
+ /// before any plan node is built (see [org.apache.pinot.core.plan.maker.InstancePlanMakerImplV2]), so this
+ /// accessor should normally only observe `ON` or `OFF`. If it is reached with the mode still `AUTO` -- meaning
+ /// the gate was reached out of order, before resolution -- it logs a warning and falls back to `OFF` rather than
+ /// failing the query. Callers that legitimately run before resolution must read
+ /// [#getSortedSelectionMergeMode()] instead.
public boolean isSortedSelectionMergeEnabled() {
- return _sortedSelectionMergeEnabled;
+ if (_sortedSelectionMergeMode == SortedSelectionMergeMode.AUTO) {
+ // This accessor is called once per segment (and again for the combine), so warn only once per query rather
+ // than flooding the log if the "should never happen" unresolved AUTO is ever actually hit.
+ if (!_unresolvedSortedSelectionMergeModeWarned) {
+ _unresolvedSortedSelectionMergeModeWarned = true;
+ LOGGER.warn("Sorted selection merge mode was read before the plan maker resolved it from AUTO; "
+ + "treating as OFF");
+ }
+ return false;
+ }
+ return _sortedSelectionMergeMode != SortedSelectionMergeMode.OFF;
+ }
+
+ public double getSortedSelectionMergeAutoMinSortedRatio() {
+ return _sortedSelectionMergeAutoMinSortedRatio;
}
- public void setSortedSelectionMergeEnabled(boolean sortedSelectionMergeEnabled) {
- _sortedSelectionMergeEnabled = sortedSelectionMergeEnabled;
+ public void setSortedSelectionMergeAutoMinSortedRatio(double sortedSelectionMergeAutoMinSortedRatio) {
+ _sortedSelectionMergeAutoMinSortedRatio = sortedSelectionMergeAutoMinSortedRatio;
}
public int getSortedSelectionMergeBlockSize() {
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/CombineSlowOperatorsTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/CombineSlowOperatorsTest.java
index 115ea33a0e5..da476ee0dbd 100644
--- a/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/CombineSlowOperatorsTest.java
+++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/CombineSlowOperatorsTest.java
@@ -170,22 +170,10 @@ public class CombineSlowOperatorsTest {
List<Operator> operators = getOperators(ready, minMaxSegmentSupplier());
QueryContext queryContext = QueryContextConverterUtils.getQueryContext("SELECT * FROM testTable ORDER BY column");
queryContext.setEndTimeMs(System.currentTimeMillis() + 10000);
- // Single-stage mode: nextBlock() drives the whole merge synchronously on the (cancellable) caller thread.
+ // getNextBlock() drives the merge synchronously on the (cancellable) caller thread, so the interrupt-on-cancel
+ // path applies to the very first call and must surface an ExceptionResultsBlock.
StreamingSelectionOrderByCombineOperator combineOperator =
- new StreamingSelectionOrderByCombineOperator(operators, queryContext, _executorService, false);
- testCancelCombineOperator(combineOperator, ready, operators);
- }
-
- @Test
- public void testCancelStreamingSelectionOrderByCombineOperatorStreamingMode() {
- CountDownLatch ready = new CountDownLatch(1);
- List<Operator> operators = getOperators(ready, minMaxSegmentSupplier());
- QueryContext queryContext = QueryContextConverterUtils.getQueryContext("SELECT * FROM testTable ORDER BY column");
- queryContext.setEndTimeMs(System.currentTimeMillis() + 10000);
- // Streaming (MSE-leaf) mode: the first getNextBlock() still drives the merge on the caller thread, so the same
- // interrupt-on-cancel path applies and must surface an ExceptionResultsBlock.
- StreamingSelectionOrderByCombineOperator combineOperator =
- new StreamingSelectionOrderByCombineOperator(operators, queryContext, _executorService, true);
+ new StreamingSelectionOrderByCombineOperator(operators, queryContext, _executorService);
testCancelCombineOperator(combineOperator, ready, operators);
}
@@ -195,11 +183,14 @@ public class CombineSlowOperatorsTest {
/// from passing vacuously if the timeout came from somewhere else.
@Test
public void testStreamingSelectionOrderByCombineOperatorHonorsDeadline() {
- List<Operator> operators = getOperators(null, minMaxSegmentSupplier());
+ // A real latch, never awaited: this test asserts that no child is driven at all, so there is nothing to wait
+ // for. It is passed only so every getOperators() call site in this class looks alike; SlowOperator null-guards
+ // the latch, so behaviour is the same either way.
+ List<Operator> operators = getOperators(new CountDownLatch(1), minMaxSegmentSupplier());
QueryContext queryContext = QueryContextConverterUtils.getQueryContext("SELECT * FROM testTable ORDER BY column");
queryContext.setEndTimeMs(System.currentTimeMillis() - 1);
StreamingSelectionOrderByCombineOperator combineOperator =
- new StreamingSelectionOrderByCombineOperator(operators, queryContext, _executorService, false);
+ new StreamingSelectionOrderByCombineOperator(operators, queryContext, _executorService);
try (QueryThreadContext ignore = QueryThreadContext.openForSseTest()) {
BaseResultsBlock resultsBlock = combineOperator.nextBlock();
assertTrue(resultsBlock instanceof ExceptionResultsBlock,
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java
index e2d9c9f8a9b..0ba5d5eb5bf 100644
--- a/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java
+++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java
@@ -29,11 +29,11 @@ import java.util.stream.Collectors;
import org.apache.commons.io.FileUtils;
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.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.plan.CombinePlanNode;
import org.apache.pinot.core.plan.PlanNode;
import org.apache.pinot.core.plan.maker.InstancePlanMakerImplV2;
@@ -55,6 +55,7 @@ import org.apache.pinot.spi.data.FieldSpec;
import org.apache.pinot.spi.data.Schema;
import org.apache.pinot.spi.data.readers.GenericRow;
import org.apache.pinot.spi.utils.CommonConstants.Server;
+import org.apache.pinot.spi.utils.CommonConstants.Server.SortedSelectionMergeMode;
import org.apache.pinot.spi.utils.ReadMode;
import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
import org.intellij.lang.annotations.Language;
@@ -63,8 +64,11 @@ import org.testng.annotations.BeforeClass;
import org.testng.annotations.Test;
import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertNull;
import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.expectThrows;
/// Combine-level tests for {@link StreamingSelectionOrderByCombineOperator} (step-3 operator) and its wiring into
@@ -73,7 +77,7 @@ import static org.testng.Assert.assertTrue;
/// <p>The streaming combine must return the same globally-sorted top-K rows as the default
/// {@link MinMaxValueBasedSelectionOrderByCombineOperator}, only (in streaming mode) spread across several bounded
/// blocks. Each functional test therefore asserts <b>streaming-vs-non-streaming parity</b>: it runs the identical query
-/// twice over the same in-memory segments - once with {@code sortedSelectionMergeEnabled=true} (asserting the new
+/// twice over the same in-memory segments - once with {@code sortedSelectionMergeMode=ON} (asserting the new
/// operator was actually selected) and once with the hint off (asserting the {@code MinMax} operator was selected) -
/// then checks the two row sets are equal as a multiset and that the streaming output is fully sorted by the order-by
/// comparator.
@@ -128,10 +132,22 @@ public class StreamingSelectionOrderByCombineOperatorTest {
/// A mix of sorted (streaming child) and physically-unsorted (single materialized top-K block child) segments,
/// exercising both SegmentCursor backings in one merge.
private List<IndexSegment> _mixedSegments;
+ /// Physically sorted on the leading order-by column, but that column also carries real nulls. `isSorted()` on the
+ /// metadata still says true, so only the null-bitmap check separates these from [#_sortedSegments].
+ private List<IndexSegment> _nullBearingSortedSegments;
+ /// No segment is physically sorted on the leading order-by column, so every child is a materializing top-K operator.
+ /// This is the shape AUTO must refuse and the shape ON must still get right.
+ private List<IndexSegment> _unsortedSegments;
/// Very low cardinality primary column (4 distinct values across 100 rows) so each value is a long run: exercises the
/// run/heap path in the streaming children and ties on the primary key at the prune boundary.
private List<IndexSegment> _lowCardSegments;
+ /// A sparse segment plus one beginning inside its first gap; see [#buildGapFillingSortedRecords].
+ private List<IndexSegment> _interleavedSegments;
+
+ /// Two segments whose rows tie on every order-by expression; see [#buildTiedSortedRecords].
+ private List<IndexSegment> _tiedSegments;
+
@BeforeClass
public void setUp()
throws Exception {
@@ -147,6 +163,19 @@ public class StreamingSelectionOrderByCombineOperatorTest {
_disjointSegments.add(buildSegment(SORTED_TABLE_CONFIG, "disjoint_" + i, buildDisjointSortedRecords(i), false));
}
+ _unsortedSegments = new ArrayList<>(NUM_SEGMENTS);
+ for (int i = 0; i < NUM_SEGMENTS; i++) {
+ _unsortedSegments.add(buildSegment(UNSORTED_TABLE_CONFIG, "unsorted_" + i, buildUnsortedRecords(i), false));
+ }
+
+ // Nulls occupy the lowest stored values (INT null default is Integer.MIN_VALUE), so sortedCol still ascends and
+ // the segment is still built as physically sorted -- exactly the shape where isSorted() alone is not trustworthy.
+ _nullBearingSortedSegments = new ArrayList<>(NUM_SEGMENTS);
+ for (int i = 0; i < NUM_SEGMENTS; i++) {
+ _nullBearingSortedSegments.add(
+ buildSegment(SORTED_TABLE_CONFIG, "nullSorted_" + i, buildNullBearingSortedRecords(i), true));
+ }
+
_lowCardSegments = new ArrayList<>(NUM_SEGMENTS);
for (int i = 0; i < NUM_SEGMENTS; i++) {
_lowCardSegments.add(buildSegment(SORTED_TABLE_CONFIG, "lowCard_" + i, buildLowCardinalityRecords(i), false));
@@ -158,6 +187,81 @@ public class StreamingSelectionOrderByCombineOperatorTest {
_mixedSegments.add(buildSegment(SORTED_TABLE_CONFIG, "mixedSorted_1", buildOverlappingSortedRecords(1), false));
_mixedSegments.add(buildSegment(UNSORTED_TABLE_CONFIG, "mixedUnsorted_0", buildUnsortedRecords(2), false));
_mixedSegments.add(buildSegment(UNSORTED_TABLE_CONFIG, "mixedUnsorted_1", buildUnsortedRecords(3), false));
+
+ // Two segments whose rows are identical on the full order-by key; see #testFullOrderByKeyTiesAcrossSegments.
+ _tiedSegments = new ArrayList<>(2);
+ _tiedSegments.add(buildSegment(SORTED_TABLE_CONFIG, "tied_0", buildTiedSortedRecords("a"), false));
+ _tiedSegments.add(buildSegment(SORTED_TABLE_CONFIG, "tied_1", buildTiedSortedRecords("b"), false));
+
+ // Sparse segment plus a dense one starting inside its first gap; see #testLateActivatedCursorWinsImmediately.
+ _interleavedSegments = new ArrayList<>(2);
+ _interleavedSegments.add(buildSegment(SORTED_TABLE_CONFIG, "sparse_0", buildSparseSortedRecords(), false));
+ _interleavedSegments.add(buildSegment(SORTED_TABLE_CONFIG, "denseGap_0", buildGapFillingSortedRecords(), false));
+ }
+
+ /// Leading order-by values `0, 1000, 2000, ...`: gaps wide enough for another segment's entire range to sit between
+ /// two consecutive rows, with no value shared with it.
+ private static List<GenericRow> buildSparseSortedRecords() {
+ List<GenericRow> records = new ArrayList<>(NUM_RECORDS_PER_SEGMENT);
+ for (int i = 0; i < NUM_RECORDS_PER_SEGMENT; i++) {
+ GenericRow record = new GenericRow();
+ record.putValue(SORTED_COL, i * 1000);
+ record.putValue(TAIL_COL, i);
+ // Globally unique against the gap-filling segment below, so the two never tie on the full order-by key.
+ record.putValue(VAL_COL, 2_000_000 + i);
+ record.putValue(NULLABLE_COL, i);
+ record.putValue(LONG_COL, 10_000_000_000L + i);
+ record.putValue(STR_COL, "sp_" + i);
+ records.add(record);
+ }
+ return records;
+ }
+
+ /// Rows that are identical on both order-by expressions across every segment built from this, differing only in a
+ /// non-order-by column, so the merge's tie branch is reached on every single row.
+ private static List<GenericRow> buildTiedSortedRecords(String tag) {
+ List<GenericRow> records = new ArrayList<>(NUM_RECORDS_PER_SEGMENT);
+ for (int i = 0; i < NUM_RECORDS_PER_SEGMENT; i++) {
+ GenericRow record = new GenericRow();
+ record.putValue(SORTED_COL, i);
+ record.putValue(TAIL_COL, i);
+ // Deliberately NOT unique across segments: sortedCol and valCol are the order-by key, so every row of one
+ // segment ties exactly with a row of the other.
+ record.putValue(VAL_COL, i);
+ record.putValue(NULLABLE_COL, i);
+ record.putValue(LONG_COL, 10_000_000_000L + i);
+ record.putValue(STR_COL, tag + "_" + i);
+ records.add(record);
+ }
+ return records;
+ }
+
+ /// Leading order-by values `5, 6, ..., 104`: the whole range sits inside the sparse segment's first gap (between its
+ /// rows 0 and 1000), so this segment is pruned at the start and activated only once the frontier passes 5.
+ private static List<GenericRow> buildGapFillingSortedRecords() {
+ List<GenericRow> records = new ArrayList<>(NUM_RECORDS_PER_SEGMENT);
+ for (int i = 0; i < NUM_RECORDS_PER_SEGMENT; i++) {
+ GenericRow record = new GenericRow();
+ record.putValue(SORTED_COL, 5 + i);
+ record.putValue(TAIL_COL, i);
+ record.putValue(VAL_COL, 3_000_000 + i);
+ record.putValue(NULLABLE_COL, i);
+ record.putValue(LONG_COL, 20_000_000_000L + i);
+ record.putValue(STR_COL, "gf_" + i);
+ records.add(record);
+ }
+ return records;
+ }
+
+ /// Like [#buildOverlappingSortedRecords] but the leading order-by column itself is null in the first few rows.
+ private static List<GenericRow> buildNullBearingSortedRecords(int index) {
+ List<GenericRow> records = buildOverlappingSortedRecords(index);
+ for (int i = 0; i < 3; i++) {
+ GenericRow record = records.get(i);
+ record.putValue(SORTED_COL, null);
+ record.addNullValueField(SORTED_COL);
+ }
+ return records;
}
private static List<GenericRow> buildOverlappingSortedRecords(int index) {
@@ -264,8 +368,9 @@ public class StreamingSelectionOrderByCombineOperatorTest {
@Test
public void testDescIncompatibleFallbackParity() {
- // allowReverseOrder=false + DESC -> the streaming child cannot scan descending, so SelectionPlanNode emits the
- // materialized DESC top-K block; the combine still routes to the streaming combine and merges single-block cursors.
+ // ON forces the path unconditionally, so allowReverseOrder=false + DESC still reaches the streaming combine over
+ // materialized DESC top-K blocks -- slower than the baseline, but it must stay correct. AUTO refuses this shape
+ // (see testAutoDoesNotSelectStreamingForDescWithoutReverseOrder).
assertParity(_sortedSegments,
"SET allowReverseOrder=false; SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol DESC, valCol DESC "
+ "LIMIT 50", false);
@@ -340,7 +445,7 @@ public class StreamingSelectionOrderByCombineOperatorTest {
// Disjoint, globally-increasing ranges + small LIMIT: only the lowest segment can contribute, the rest are pruned
// (never acquired or scanned), so far fewer than all docs are scanned.
Result result = run(_disjointSegments, "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 5",
- true, false, false, 0);
+ true, false, true, 0);
assertTrue(result._combineOperator instanceof StreamingSelectionOrderByCombineOperator);
assertEquals(result._rows.size(), 5);
for (int i = 0; i < 5; i++) {
@@ -353,56 +458,213 @@ public class StreamingSelectionOrderByCombineOperatorTest {
"Only the lowest-range segment should be scanned, but docs scanned was: " + result._numDocsScanned);
}
+ /// A cursor activated *late* can hold a smaller row than the retained leader. Two things must be right for it to
+ /// land in the correct place: the leader is re-compared after `activateEligibleCursors()` rather than before, and
+ /// the pruning frontier is the row about to be emitted rather than the last one emitted. Either mistake emits the
+ /// leader's larger row first.
+ ///
+ /// The sparse segment runs `0, 1000, 2000, ...` and the gap filler `5, 6, ..., 104`, so the filler is pruned while
+ /// the frontier is 0 and becomes eligible only when the leader reaches 1000 -- where its head (5) must displace it at
+ /// once. The ranges share no value, so the expected sequence below is a strict total order with no tie to mask a
+ /// misordering. Parity does not cover this: [#assertMultisetEquals] normalises the reordering away, hence the
+ /// positional assertion.
@Test
- public void testEmptyResultSchemaFallback() {
- // A filter that matches nothing: every streaming child returns no rows, so the combine rebuilds the result schema
- // from the segment metadata (order-by expressions first) rather than from a child block.
+ public void testLateActivatedCursorWinsImmediately() {
+ @Language("sql") String query =
+ "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 8";
+ Result result = run(_interleavedSegments, query, true, false, true, 3);
+ assertTrue(result._combineOperator instanceof StreamingSelectionOrderByCombineOperator);
+ // 0 from the sparse segment, then the gap filler's 5, 6, 7, ... -- its next row (1000) comes far later.
+ int[] expected = {0, 5, 6, 7, 8, 9, 10, 11};
+ assertEquals(result._rows.size(), expected.length);
+ for (int i = 0; i < expected.length; i++) {
+ assertEquals((int) result._rows.get(i)[0], expected[i],
+ "Unexpected value at position " + i + "; a late-activated cursor was not given the lead immediately");
+ }
+ assertParity(_interleavedSegments, query, false);
+ }
+
+ /// Pruning must survive leader retention: one cursor serving a long run leaves the heap empty, and a frontier read
+ /// from the heap alone would call that "unknown" and activate every remaining segment -- switching pruning off on
+ /// exactly the near-disjoint shape this operator exists for. [#testPruningSkipsOutOfTopKSegments] retains the leader
+ /// only a few times; this limit spans several blocks, widening the empty-heap window.
+ @Test
+ public void testPruningSurvivesALongRetainedRun() {
+ Result result = run(_disjointSegments,
+ "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 60", true, false, true, 7);
+ assertTrue(result._combineOperator instanceof StreamingSelectionOrderByCombineOperator);
+ assertEquals(result._rows.size(), 60);
+ assertTrue(result._numBlocks > 1, "Fixture must span several blocks, else the retained run is too short");
+ for (int i = 0; i < 60; i++) {
+ assertEquals((int) result._rows.get(i)[0], i, "Unexpected value at position " + i);
+ }
+ assertTrue(result._numDocsScanned <= NUM_RECORDS_PER_SEGMENT,
+ "All 60 rows come from the lowest-range segment, so no other segment may be activated; docs scanned was: "
+ + result._numDocsScanned);
+ }
+
+ /// The merge cedes the lead only on a strict loss, so a tie keeps the incumbent -- a branch no other fixture reaches,
+ /// because they all give `valCol` a globally unique value precisely to avoid ties. Here every row of one segment ties
+ /// with a row of the other on the complete order-by key.
+ ///
+ /// What is asserted is that ties are *harmless*: nothing is dropped, duplicated or emitted out of order. Which of two
+ /// fully-tied rows comes first is deliberately not pinned -- the incumbent-wins rule is an artefact of leader
+ /// retention, the heap's own tie-breaking was arbitrary before it, and neither path promises stability. Asserting it
+ /// would freeze behaviour the operator does not guarantee.
+ /// The DESC counterpart of [#testPruningSurvivesALongRetainedRun]. The pruning test in `sortsBeyond` flips its
+ /// comparison on `_asc`, and every other live-pruning test runs ASC only -- so a sign flip or swapped operand in the
+ /// DESC half would otherwise go unnoticed, silently either over-scanning or, worse, pruning a segment that still had
+ /// rows to give.
+ @Test
+ public void testPruningSurvivesALongRetainedRunDesc() {
+ Result result = run(_disjointSegments,
+ "SET allowReverseOrder=true; SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol DESC, valCol DESC "
+ + "LIMIT 60", true, false, true, 7);
+ assertTrue(result._combineOperator instanceof StreamingSelectionOrderByCombineOperator);
+ assertEquals(result._rows.size(), 60);
+ assertTrue(result._numBlocks > 1, "Fixture must span several blocks, else the retained run is too short");
+ // DESC activates the highest-range segment first; its top value counts down from there.
+ int top = (NUM_SEGMENTS - 1) * 1000 + NUM_RECORDS_PER_SEGMENT - 1;
+ for (int i = 0; i < 60; i++) {
+ assertEquals((int) result._rows.get(i)[0], top - i, "Unexpected value at position " + i);
+ }
+ assertTrue(result._numDocsScanned <= NUM_RECORDS_PER_SEGMENT,
+ "All 60 rows come from the highest-range segment, so no other segment may be activated; docs scanned was: "
+ + result._numDocsScanned);
+ }
+
+ @Test
+ public void testFullOrderByKeyTiesAcrossSegments() {
+ // Explicit limit: the default of 10 would stop the merge long before most ties are reached.
+ @Language("sql") String query =
+ "SELECT sortedCol, valCol, strCol FROM testTable ORDER BY sortedCol, valCol LIMIT 200";
+ assertParity(_tiedSegments, query, false);
+
+ Result streamed = run(_tiedSegments, query, true, false, true, 5);
+ assertEquals(streamed._rows.size(), 2 * NUM_RECORDS_PER_SEGMENT, "Every tied row must survive the merge");
+ for (int i = 0; i < streamed._rows.size(); i++) {
+ // Both segments contribute one row per sortedCol value, so the merged sequence is each value twice.
+ assertEquals((int) streamed._rows.get(i)[0], i / 2, "Unexpected value at position " + i);
+ }
+ }
+
+ @Test
+ public void testEmptyResultOnStreamingPath() {
+ // A filter that matches nothing. Every streaming child still emits one empty schema-carrying block, so the combine
+ // does capture a DataSchema -- but no rows ever reach the output buffer, so it never flushes a data block and
+ // exits through the terminal MetadataResultsBlock, whose schema is unconditionally null.
+ // Carrying the schema end-to-end therefore means changing that terminal-block convention, which is shared by every
+ // BaseStreamingCombineOperator subclass; both the MSE leaf and the gRPC streaming reduce tolerate the null today,
+ // so it is pinned here rather than left free to change unnoticed.
+ // TODO: revisit if the terminal block ever carries a schema -- the per-child guarantee that would need is now in
+ // place. See https://github.com/apache/pinot/pull/19120#discussion_r3871714002
Result result = run(_sortedSegments,
"SELECT sortedCol, valCol FROM testTable WHERE sortedCol < 0 ORDER BY sortedCol, valCol LIMIT 10", true, false,
- false, 0);
+ true, 0);
assertTrue(result._combineOperator instanceof StreamingSelectionOrderByCombineOperator);
assertTrue(result._rows.isEmpty(), "Expected an empty result, got: " + result._rows.size() + " rows");
- assertEquals(result._schema, new DataSchema(new String[]{SORTED_COL, VAL_COL},
- new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.INT}));
+ assertNull(result._schema, "A zero-match streaming query carries no schema today; see the comment above");
+ }
+
+ @Test
+ public void testZeroMatchSegmentsMergedAlongsideAMatchingSegment() {
+ // Disjoint segment i covers [i * 1000, i * 1000 + 100), so this filter matches every row of the first segment and
+ // nothing at all in the rest. Each zero-match child still emits its empty schema-carrying block, so the merge runs
+ // cursors that never set a head alongside one that does -- the shape this feature targets once pruning or
+ // colocation leaves only part of the fan-out in range.
+ @Language("sql") String query =
+ "SELECT sortedCol, valCol FROM testTable WHERE sortedCol < 1000 ORDER BY sortedCol, valCol LIMIT 5";
+ assertParity(_disjointSegments, query, false);
+ Result streamed = run(_disjointSegments, query, true, false, true, 3);
+ assertFalse(streamed._rows.isEmpty(), "The in-range segment must contribute rows, else the test is vacuous");
+ assertNotNull(streamed._schema, "A merge that produced rows must carry a schema");
+ }
+
+ /// The operator passes a null merger to the base class and replaces the base streaming loop wholesale. The two base
+ /// entry points that would dereference that merger must fail loud, so a future change routing back through them
+ /// breaks at the seam instead of with an NPE at query time.
+ @Test
+ public void testBaseWorkerEntryPointsFailLoud() {
+ Operator<?> combineOperator = planCombineOperator(_sortedSegments,
+ hintedContext("SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50"), true);
+ assertTrue(combineOperator instanceof StreamingSelectionOrderByCombineOperator,
+ "Expected the streaming combine operator, got: " + combineOperator.getClass().getSimpleName());
+ StreamingSelectionOrderByCombineOperator streamingCombine =
+ (StreamingSelectionOrderByCombineOperator) combineOperator;
+ expectThrows(IllegalStateException.class, () -> streamingCombine.isQuerySatisfied(null, null));
+ expectThrows(IllegalStateException.class, streamingCombine::processSegments);
}
@Test
public void testHintOffSelectsMinMaxOperator() {
// Default behavior is unchanged when the hint is off: the classic MinMax operator is still selected.
- Result result = run(_sortedSegments, "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50",
- false, false, false, 0);
- assertTrue(result._combineOperator instanceof MinMaxValueBasedSelectionOrderByCombineOperator,
+ QueryContext queryContext = QueryContextConverterUtils.getQueryContext(
+ "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50");
+ queryContext.setEndTimeMs(System.currentTimeMillis() + Server.DEFAULT_QUERY_EXECUTOR_TIMEOUT_MS);
+ Operator<?> combineOperator = planCombineOperator(_sortedSegments, queryContext, true);
+ assertTrue(combineOperator instanceof MinMaxValueBasedSelectionOrderByCombineOperator,
"Hint off must keep the default MinMax combine operator, got: "
- + result._combineOperator.getClass().getSimpleName());
+ + combineOperator.getClass().getSimpleName());
}
@Test
- public void testNonIdentifierOrderByFallsBack() {
- // Even with the hint on, a non-identifier first order-by expression falls back to SelectionOrderByCombineOperator
- // (the streaming operator and its segment-level counterpart only support a leading identifier).
- Result result = run(_sortedSegments,
- "SELECT sortedCol, valCol FROM testTable ORDER BY ADD(sortedCol, 1), valCol LIMIT 50", true, false, false, 0);
- assertEquals(result._combineOperator.getClass(), SelectionOrderByCombineOperator.class,
+ public void testHintOnBlockingPathKeepsMinMaxOperator() {
+ // The merge can only stream where there is a ResultsBlockStreamer. With none, the hint must not select it: the
+ // blocking caller takes a single nextBlock(), so the merge could never flush more than one block.
+ Operator<?> combineOperator = planCombineOperator(_sortedSegments,
+ hintedContext("SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50"), false);
+ assertTrue(combineOperator instanceof MinMaxValueBasedSelectionOrderByCombineOperator,
+ "Hint on the blocking path must keep the MinMax combine operator, got: "
+ + combineOperator.getClass().getSimpleName());
+ }
+
+ @Test
+ public void testBlockingInstancePlanDisablesTheHint() {
+ // The leaf gate reads the same flag as the combine gate but cannot see the streamer, so the blocking instance
+ // plan clears the flag outright rather than leaving the option half-honored: streaming leaves swapped in under
+ // a blocking MinMax combine that cannot drive them.
+ QueryContext queryContext =
+ hintedContext("SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50");
+ List<SegmentContext> segmentContexts = new ArrayList<>(_sortedSegments.size());
+ for (IndexSegment segment : _sortedSegments) {
+ segmentContexts.add(new SegmentContext(segment));
+ }
+ PLAN_MAKER.makeInstancePlan(segmentContexts, queryContext, EXECUTOR);
+ assertFalse(queryContext.isSortedSelectionMergeEnabled(),
+ "The blocking instance plan must disable the streaming selection merge for the whole plan");
+ // The observable consequence: leaves planned for this query are the materializing operators, not the multi-block
+ // streaming one that only the streaming combine knows how to drain.
+ Operator<?> leafOperator = PLAN_MAKER.makeSegmentPlanNode(segmentContexts.get(0), queryContext).run();
+ assertFalse(leafOperator instanceof StreamingSelectionOrderByOperator,
+ "The blocking path must not plan a multi-block streaming leaf, got: "
+ + leafOperator.getClass().getSimpleName());
+ }
+
+ @Test
+ public void testNonIdentifierOrderByFallsBackWithTheModeOff() {
+ // Default behaviour for a non-identifier first order-by is unchanged: SelectionOrderByCombineOperator. Only an
+ // explicit ON reaches the streaming merge for this shape (see testOnForcesStreamingForAnExpressionOrderBy);
+ // AUTO refuses it (see testAutoDoesNotSelectStreamingForAnExpressionOrderBy).
+ QueryContext queryContext = QueryContextConverterUtils.getQueryContext(
+ "SELECT sortedCol, valCol FROM testTable ORDER BY ADD(sortedCol, 1), valCol LIMIT 50");
+ queryContext.setEndTimeMs(System.currentTimeMillis() + Server.DEFAULT_QUERY_EXECUTOR_TIMEOUT_MS);
+ Operator<?> combineOperator = planCombineOperator(_sortedSegments, queryContext, true);
+ assertEquals(combineOperator.getClass(), SelectionOrderByCombineOperator.class,
"Non-identifier first order-by must fall back to SelectionOrderByCombineOperator, got: "
- + result._combineOperator.getClass().getSimpleName());
+ + combineOperator.getClass().getSimpleName());
}
- /// Asserts streaming-vs-non-streaming parity for {@code query}. Runs the MinMax combine (hint off) as the reference,
- /// then runs the streaming combine in BOTH single-block mode and bounded multi-block streaming mode, asserting each
- /// selects the streaming operator and produces rows that are sorted by the order-by comparator and equal the MinMax
- /// rows as a multiset. The streaming variant additionally checks the bounded-flush invariants.
+ /// Asserts streaming-vs-blocking parity for {@code query}. Runs the MinMax combine on the blocking path (hint off)
+ /// as the reference, then runs the streaming combine on the streaming path, asserting it selects the streaming
+ /// operator and produces rows that are sorted by the order-by comparator and equal the MinMax rows as a multiset,
+ /// plus the bounded-flush invariants. A small block size forces several bounded data blocks before the metadata
+ /// block, genuinely exercising the streaming flush path rather than a single trimmed block.
private void assertParity(List<IndexSegment> segments, @Language("sql") String query, boolean nullHandling) {
Result baseline = run(segments, query, false, nullHandling, false, 0);
assertEquals(baseline._combineOperator.getClass(), MinMaxValueBasedSelectionOrderByCombineOperator.class,
"Baseline must be the MinMax combine operator, got: " + baseline._combineOperator.getClass().getSimpleName());
Comparator<Object[]> comparator = orderByComparator(query, nullHandling);
- // Classic single-stage path (null streamer): the streaming combine flushes the whole merge as one block.
- Result singleStage = run(segments, query, true, nullHandling, false, 0);
- assertStreamingParity(singleStage, baseline, comparator, query, false, 0);
-
- // MSE leaf path (non-null streamer): a small block size forces several bounded data blocks before the metadata
- // block, genuinely exercising the streaming flush path rather than a single trimmed block.
int blockSize = 3;
Result streamed = run(segments, query, true, nullHandling, true, blockSize);
assertStreamingParity(streamed, baseline, comparator, query, true, blockSize);
@@ -429,31 +691,266 @@ public class StreamingSelectionOrderByCombineOperatorTest {
}
}
+ /// Plans the combine operator that would run {@code queryContext} over {@code segments}, without driving it. Use
+ /// this for gate assertions: an operator that is not the streaming combine (the fallbacks) does not terminate with
+ /// a metadata block, so it cannot be driven by the streaming loop in [#run].
+ private static Operator<?> planCombineOperator(List<IndexSegment> segments, QueryContext queryContext,
+ boolean streaming) {
+ List<PlanNode> planNodes = new ArrayList<>(segments.size());
+ for (IndexSegment segment : segments) {
+ SegmentContext segmentContext = new SegmentContext(segment);
+ planNodes.add(streaming ? PLAN_MAKER.makeStreamingSegmentPlanNode(segmentContext, queryContext)
+ : PLAN_MAKER.makeSegmentPlanNode(segmentContext, queryContext));
+ }
+ ResultsBlockStreamer streamer = streaming ? block -> {
+ } : null;
+ return new CombinePlanNode(planNodes, queryContext, EXECUTOR, streamer).run();
+ }
+
+ @Test
+ public void testAutoSelectsStreamingWhenAllSegmentsAreSorted() {
+ // AUTO is resolved from segment metadata by the streaming instance plan, before any plan node is built.
+ QueryContext queryContext = modeContext(
+ "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50", SortedSelectionMergeMode.AUTO);
+ makeStreamingInstancePlan(_sortedSegments, queryContext);
+ assertEquals(queryContext.getSortedSelectionMergeMode(), SortedSelectionMergeMode.ON,
+ "AUTO must resolve to ON when every segment is sorted on the leading order-by column");
+ assertTrue(planCombineOperator(_sortedSegments, queryContext, true)
+ instanceof StreamingSelectionOrderByCombineOperator,
+ "The resolved mode must select the streaming combine");
+ }
+
+ @Test
+ public void testAutoSelectsMinMaxWhenNoSegmentIsSorted() {
+ // Every child would be a materializing SelectionOrderByOperator doing a full scan plus top-K synchronously on the
+ // consumer thread, one segment after another -- strictly worse than the parallel MinMax combine.
+ QueryContext queryContext = modeContext(
+ "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50", SortedSelectionMergeMode.AUTO);
+ makeStreamingInstancePlan(_unsortedSegments, queryContext);
+ assertEquals(queryContext.getSortedSelectionMergeMode(), SortedSelectionMergeMode.OFF,
+ "AUTO must resolve to OFF when no segment is sorted on the leading order-by column");
+ assertTrue(planCombineOperator(_unsortedSegments, queryContext, true)
+ instanceof MinMaxValueBasedSelectionOrderByCombineOperator,
+ "The resolved mode must keep the MinMax combine");
+ }
+
+ @Test
+ public void testAutoHonoursTheMinSortedRatioThreshold() {
+ // _mixedSegments is 2 sorted of 4, so a ratio of exactly 0.5 must pass and anything above it must fail. This pins
+ // the comparison as >= rather than >, and pins that the threshold is read from the query context.
+ String query = "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50";
+
+ QueryContext atThreshold = modeContext(query, SortedSelectionMergeMode.AUTO);
+ atThreshold.setSortedSelectionMergeAutoMinSortedRatio(0.5);
+ makeStreamingInstancePlan(_mixedSegments, atThreshold);
+ assertEquals(atThreshold.getSortedSelectionMergeMode(), SortedSelectionMergeMode.ON,
+ "A sorted ratio equal to the threshold must select the streaming merge");
+
+ QueryContext aboveThreshold = modeContext(query, SortedSelectionMergeMode.AUTO);
+ aboveThreshold.setSortedSelectionMergeAutoMinSortedRatio(0.75);
+ makeStreamingInstancePlan(_mixedSegments, aboveThreshold);
+ assertEquals(aboveThreshold.getSortedSelectionMergeMode(), SortedSelectionMergeMode.OFF,
+ "A sorted ratio below the threshold must keep the MinMax combine");
+ }
+
+ @Test
+ public void testAutoIgnoresNullsInTheLeadingColumn() {
+ // Pins the known gap documented by the TODO in InstancePlanMakerImplV2#isSortedEnoughForStreamingMerge, so it is
+ // enforced by CI rather than only described. AUTO resolves from physical sortedness alone, because the null check
+ // that SelectionPlanNode#isColumnSorted adds reads the segment's mapped buffer and AUTO runs before any acquire.
+ // So a null-bearing sorted column resolves to ON with null handling either off or on -- but with it on, the leaf
+ // gate then refuses to stream, and the merge runs over fully materialized children. Correct, just not fast.
+ //
+ // Tighten this to OFF for the null-handling case when the follow-up lands.
+ // See https://github.com/apache/pinot/pull/19120#discussion_r3871713975
+ String query = "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50";
+
+ QueryContext nullHandlingOff = modeContext(query, SortedSelectionMergeMode.AUTO);
+ makeStreamingInstancePlan(_nullBearingSortedSegments, nullHandlingOff);
+ assertEquals(nullHandlingOff.getSortedSelectionMergeMode(), SortedSelectionMergeMode.ON,
+ "With null handling off the segments are sorted, so AUTO must select the streaming merge");
+
+ QueryContext nullHandlingOn = modeContext(query, SortedSelectionMergeMode.AUTO);
+ nullHandlingOn.setNullHandlingEnabled(true);
+ makeStreamingInstancePlan(_nullBearingSortedSegments, nullHandlingOn);
+ assertEquals(nullHandlingOn.getSortedSelectionMergeMode(), SortedSelectionMergeMode.ON,
+ "AUTO currently reads physical sortedness only, so null handling does not change the decision");
+ // The leaf, which runs post-acquire and can afford the null check, still refuses to stream these segments.
+ List<SegmentContext> segmentContexts = new ArrayList<>(1);
+ segmentContexts.add(new SegmentContext(_nullBearingSortedSegments.get(0)));
+ Operator<?> leafOperator = PLAN_MAKER.makeSegmentPlanNode(segmentContexts.get(0), nullHandlingOn).run();
+ assertFalse(leafOperator instanceof StreamingSelectionOrderByOperator,
+ "A null-bearing leading column must not produce a streaming leaf under null handling, got: "
+ + leafOperator.getClass().getSimpleName());
+ }
+
+ @Test
+ public void testAutoSelectsMinMaxWhenThereAreNoSegments() {
+ QueryContext queryContext = modeContext(
+ "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50", SortedSelectionMergeMode.AUTO);
+ makeStreamingInstancePlan(List.of(), queryContext);
+ assertEquals(queryContext.getSortedSelectionMergeMode(), SortedSelectionMergeMode.OFF,
+ "An empty segment set must resolve to OFF rather than dividing by zero");
+ }
+
+ @Test
+ public void testUnresolvedAutoFallsBackToOff() {
+ // AUTO is a request for a decision, not a decision, but a gate reached before the plan maker resolved it must
+ // not fail the query either. isSortedSelectionMergeEnabled() degrades an unresolved AUTO to OFF (with a
+ // warning) rather than forcing the merge unconditionally and silently ignoring the ratio, so the gate should
+ // select the same fallback operator it would for OFF.
+ QueryContext queryContext = modeContext(
+ "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50", SortedSelectionMergeMode.AUTO);
+ Operator<?> combineOperator = planCombineOperator(_sortedSegments, queryContext, true);
+ assertTrue(combineOperator instanceof MinMaxValueBasedSelectionOrderByCombineOperator,
+ "An unresolved AUTO must degrade to OFF, got: " + combineOperator.getClass().getSimpleName());
+ }
+
+ @Test
+ public void testAutoDoesNotSelectStreamingForAnExpressionOrderBy() {
+ // Sortedness is only defined for a physical column; forcing the merge over an expression is what ON is for.
+ QueryContext queryContext =
+ modeContext("SELECT sortedCol, valCol FROM testTable ORDER BY ADD(sortedCol, 1) LIMIT 50",
+ SortedSelectionMergeMode.AUTO);
+ makeStreamingInstancePlan(_sortedSegments, queryContext);
+ assertEquals(queryContext.getSortedSelectionMergeMode(), SortedSelectionMergeMode.OFF,
+ "AUTO must not select the streaming merge for a non-identifier leading order-by expression");
+ }
+
+ @Test
+ public void testAutoDoesNotSelectStreamingForDescWithoutReverseOrder() {
+ // DataSourceMetadata.isSorted() reports ascending physical order, so with allowReverseOrder off no leaf can scan
+ // sortedCol descending: every segment falls back to the materialized SelectionPartiallyOrderedByDescOperation.
+ // Resolving to ON here would install the streaming combine over children that each materialize their top-K
+ // serially, losing the MinMax combine's parallelism and its min/max segment pruning for no gain.
+ QueryContext queryContext =
+ modeContext("SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol DESC, valCol DESC LIMIT 50",
+ SortedSelectionMergeMode.AUTO);
+ makeStreamingInstancePlan(_sortedSegments, queryContext);
+ assertEquals(queryContext.getSortedSelectionMergeMode(), SortedSelectionMergeMode.OFF,
+ "AUTO must not select the streaming merge for a DESC order-by that no leaf can stream");
+ assertTrue(planCombineOperator(_sortedSegments, queryContext, true)
+ instanceof MinMaxValueBasedSelectionOrderByCombineOperator,
+ "The resolved mode must keep the MinMax combine");
+ }
+
+ @Test
+ public void testAutoSelectsStreamingForDescWithReverseOrder() {
+ // With allowReverseOrder on, getSortedByProject() hands the leaf a DESC-compatible project operator, so the
+ // streaming children this merge needs do get built and AUTO should take the path.
+ QueryContext queryContext = modeContext(
+ "SET allowReverseOrder=true; SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol DESC, valCol DESC "
+ + "LIMIT 50", SortedSelectionMergeMode.AUTO);
+ makeStreamingInstancePlan(_sortedSegments, queryContext);
+ assertEquals(queryContext.getSortedSelectionMergeMode(), SortedSelectionMergeMode.ON,
+ "AUTO must select the streaming merge for DESC once reverse iteration is allowed");
+ assertTrue(planCombineOperator(_sortedSegments, queryContext, true)
+ instanceof StreamingSelectionOrderByCombineOperator,
+ "The resolved mode must select the streaming combine");
+ }
+
+ @Test
+ public void testAutoHonoursTheMinSortedRatioThresholdForDesc() {
+ // The DESC gate is a precondition, not a replacement for the ratio check: once reverse iteration is allowed, a
+ // DESC query must still clear the same threshold an ASC one does. _mixedSegments is 2 sorted of 4, so 0.5 passes
+ // and 0.75 does not, exactly as in testAutoHonoursTheMinSortedRatioThreshold.
+ String query = "SET allowReverseOrder=true; SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol DESC, "
+ + "valCol DESC LIMIT 50";
+
+ QueryContext atThreshold = modeContext(query, SortedSelectionMergeMode.AUTO);
+ atThreshold.setSortedSelectionMergeAutoMinSortedRatio(0.5);
+ makeStreamingInstancePlan(_mixedSegments, atThreshold);
+ assertEquals(atThreshold.getSortedSelectionMergeMode(), SortedSelectionMergeMode.ON,
+ "A DESC sorted ratio equal to the threshold must select the streaming merge");
+
+ QueryContext aboveThreshold = modeContext(query, SortedSelectionMergeMode.AUTO);
+ aboveThreshold.setSortedSelectionMergeAutoMinSortedRatio(0.75);
+ makeStreamingInstancePlan(_mixedSegments, aboveThreshold);
+ assertEquals(aboveThreshold.getSortedSelectionMergeMode(), SortedSelectionMergeMode.OFF,
+ "A DESC sorted ratio below the threshold must keep the MinMax combine");
+ }
+
+ @Test
+ public void testAutoOnlyChecksTheLeadingOrderByDirection() {
+ // Only the leading expression rides the segment's physical order, so a DESC tail is sorted in memory and needs no
+ // reverse scan. Pins the gate as "leading expression is DESC", not "any expression is DESC".
+ QueryContext queryContext =
+ modeContext("SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol DESC LIMIT 50",
+ SortedSelectionMergeMode.AUTO);
+ makeStreamingInstancePlan(_sortedSegments, queryContext);
+ assertEquals(queryContext.getSortedSelectionMergeMode(), SortedSelectionMergeMode.ON,
+ "An ASC leading expression must select the streaming merge regardless of the tail's direction");
+ assertTrue(planCombineOperator(_sortedSegments, queryContext, true)
+ instanceof StreamingSelectionOrderByCombineOperator,
+ "The resolved mode must select the streaming combine");
+ }
+
+ @Test
+ public void testOnForcesStreamingForAnExpressionOrderBy() {
+ // ON takes the path unconditionally, expressions included, so it can be benchmarked against the default. There is
+ // no leading physical column, so frontier pruning is inert and every cursor is activated in plan order.
+ QueryContext queryContext =
+ modeContext("SELECT sortedCol, valCol FROM testTable ORDER BY ADD(sortedCol, 1) LIMIT 50",
+ SortedSelectionMergeMode.ON);
+ makeStreamingInstancePlan(_sortedSegments, queryContext);
+ assertEquals(queryContext.getSortedSelectionMergeMode(), SortedSelectionMergeMode.ON,
+ "ON must pass through the AUTO resolution untouched");
+ assertTrue(planCombineOperator(_sortedSegments, queryContext, true)
+ instanceof StreamingSelectionOrderByCombineOperator,
+ "ON must select the streaming combine even for an expression order-by");
+ }
+
+ @Test
+ public void testOnOverUnsortedSegmentsMatchesTheBaseline() {
+ // Forced over unsorted segments the merge runs on materialized SelectionOrderByOperator children. Slower, but it
+ // must still produce exactly the baseline result -- that is what makes it usable for A/B.
+ assertParity(_unsortedSegments, "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50",
+ false);
+ }
+
+ /// Runs the streaming instance plan for its side effect on {@code queryContext}: resolving AUTO from segment
+ /// metadata. The plan itself is discarded; the gate tests assert on the resolved mode and on the operator that
+ /// [#planCombineOperator] then selects.
+ private static void makeStreamingInstancePlan(List<IndexSegment> segments, QueryContext queryContext) {
+ List<SegmentContext> segmentContexts = new ArrayList<>(segments.size());
+ for (IndexSegment segment : segments) {
+ segmentContexts.add(new SegmentContext(segment));
+ }
+ PLAN_MAKER.makeStreamingInstancePlan(segmentContexts, queryContext, EXECUTOR, block -> {
+ });
+ }
+
+ /// Builds a query context for {@code query} with the streaming merge mode set to {@code mode}.
+ private static QueryContext modeContext(@Language("sql") String query, SortedSelectionMergeMode mode) {
+ QueryContext queryContext = QueryContextConverterUtils.getQueryContext(query);
+ queryContext.setSortedSelectionMergeMode(mode);
+ queryContext.setEndTimeMs(System.currentTimeMillis() + Server.DEFAULT_QUERY_EXECUTOR_TIMEOUT_MS);
+ return queryContext;
+ }
+
+ /// Builds a query context for {@code query} with the streaming merge hint on.
+ private static QueryContext hintedContext(@Language("sql") String query) {
+ QueryContext queryContext = QueryContextConverterUtils.getQueryContext(query);
+ queryContext.setSortedSelectionMergeMode(SortedSelectionMergeMode.ON);
+ queryContext.setEndTimeMs(System.currentTimeMillis() + Server.DEFAULT_QUERY_EXECUTOR_TIMEOUT_MS);
+ return queryContext;
+ }
+
/// Runs one combine over {@code segments} and collects its rows, blocks, schema and docs-scanned stat.
private Result run(List<IndexSegment> segments, @Language("sql") String query, boolean hintOn, boolean nullHandling,
boolean streaming, int blockSize) {
QueryContext queryContext = QueryContextConverterUtils.getQueryContext(query);
queryContext.setNullHandlingEnabled(nullHandling);
if (hintOn) {
- queryContext.setSortedSelectionMergeEnabled(true);
+ queryContext.setSortedSelectionMergeMode(SortedSelectionMergeMode.ON);
if (blockSize > 0) {
queryContext.setSortedSelectionMergeBlockSize(blockSize);
}
}
queryContext.setEndTimeMs(System.currentTimeMillis() + Server.DEFAULT_QUERY_EXECUTOR_TIMEOUT_MS);
- List<PlanNode> planNodes = new ArrayList<>(segments.size());
- for (IndexSegment segment : segments) {
- SegmentContext segmentContext = new SegmentContext(segment);
- planNodes.add(streaming ? PLAN_MAKER.makeStreamingSegmentPlanNode(segmentContext, queryContext)
- : PLAN_MAKER.makeSegmentPlanNode(segmentContext, queryContext));
- }
- ResultsBlockStreamer streamer = streaming ? block -> {
- } : null;
- CombinePlanNode combinePlanNode = new CombinePlanNode(planNodes, queryContext, EXECUTOR, streamer);
-
Result result = new Result();
- Operator<?> combineOperator = combinePlanNode.run();
+ Operator<?> combineOperator = planCombineOperator(segments, queryContext, streaming);
result._combineOperator = combineOperator;
result._rows = new ArrayList<>();
result._blockSizes = new ArrayList<>();
@@ -533,7 +1030,8 @@ public class StreamingSelectionOrderByCombineOperatorTest {
public void tearDown()
throws IOException {
EXECUTOR.shutdownNow();
- for (List<IndexSegment> segments : List.of(_sortedSegments, _disjointSegments, _mixedSegments, _lowCardSegments)) {
+ for (List<IndexSegment> segments : List.of(_sortedSegments, _disjointSegments, _mixedSegments, _lowCardSegments,
+ _interleavedSegments, _tiedSegments)) {
for (IndexSegment segment : segments) {
segment.destroy();
}
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperatorTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperatorTest.java
index 73a071f673a..f77bdc4d4b4 100644
--- a/pinot-core/src/test/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperatorTest.java
+++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperatorTest.java
@@ -20,15 +20,28 @@ package org.apache.pinot.core.operator.query;
import java.io.File;
import java.io.IOException;
+import java.util.ArrayDeque;
import java.util.ArrayList;
+import java.util.Deque;
import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
import org.apache.commons.io.FileUtils;
+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.operator.BaseProjectOperator;
+import org.apache.pinot.core.operator.ColumnContext;
+import org.apache.pinot.core.operator.ExecutionStatistics;
+import org.apache.pinot.core.operator.blocks.ValueBlock;
import org.apache.pinot.core.operator.blocks.results.SelectionResultsBlock;
+import org.apache.pinot.core.plan.ProjectPlanNode;
import org.apache.pinot.core.plan.SelectionPlanNode;
import org.apache.pinot.core.query.request.context.QueryContext;
import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils;
+import org.apache.pinot.core.query.selection.SelectionOperatorUtils;
import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader;
import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl;
import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader;
@@ -40,6 +53,7 @@ import org.apache.pinot.spi.config.table.TableType;
import org.apache.pinot.spi.data.FieldSpec;
import org.apache.pinot.spi.data.Schema;
import org.apache.pinot.spi.data.readers.GenericRow;
+import org.apache.pinot.spi.utils.CommonConstants.Server.SortedSelectionMergeMode;
import org.apache.pinot.spi.utils.ReadMode;
import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
import org.intellij.lang.annotations.Language;
@@ -47,9 +61,15 @@ import org.testng.annotations.AfterClass;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.Test;
+import static org.mockito.AdditionalAnswers.delegatesTo;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertFalse;
import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertNull;
import static org.testng.Assert.assertTrue;
@@ -57,7 +77,7 @@ import static org.testng.Assert.assertTrue;
///
/// <p>The operator emits the same globally-sorted rows as the existing materialized selection ORDER BY operators, only
/// spread across many lazily-produced blocks. Each test therefore asserts <b>stream-vs-materialized parity</b>: it
-/// drives the streaming operator through {@link SelectionPlanNode} with {@code sortedSelectionMergeEnabled=true},
+/// drives the streaming operator through {@link SelectionPlanNode} with {@code sortedSelectionMergeMode=ON},
/// concatenates
/// every {@link Operator#nextBlock()} output until {@code null}, and asserts the concatenation equals the single block
/// the materialized operator ({@link SelectionPartiallyOrderedByLinearOperator} / {@link SelectionOrderByOperator})
@@ -103,6 +123,9 @@ public class StreamingSelectionOrderByOperatorTest {
private static final int MULTI_RUN0_SIZE = 10_500;
private static final int MULTI_RUN1_SIZE = 1_000;
private static final int NUM_MULTI_RUN_RECORDS = MULTI_RUN0_SIZE + MULTI_RUN1_SIZE;
+ /// Project block size used by the empty-block injection tests only, small enough that a modest fixture spans
+ /// several blocks so the injected empty block lands mid-scan.
+ private static final int MAX_DOCS_PER_PROJECT_BLOCK = 4;
private IndexSegment _segment;
private IndexSegment _dupSegment;
@@ -326,6 +349,448 @@ public class StreamingSelectionOrderByOperatorTest {
"Expected no nulls in the output when null handling is disabled");
}
+ /// No-tail mode, single phase: {@link StreamingSelectionOrderByOperator} must skip a zero-document project block
+ /// rather than read it as end-of-segment.
+ @Test
+ public void testEmptyProjectBlockDoesNotTruncateSortedScan() {
+ assertEmptyProjectBlocksAreSkipped(_segment, "SELECT sortedCol FROM testTable ORDER BY sortedCol LIMIT 25", 1, 2,
+ 1, 0);
+ }
+
+ /// Same, two phase: this is also the only path that asks the injected block for {@code getDocIds()}.
+ @Test
+ public void testEmptyProjectBlockDoesNotTruncateTwoPhaseSortedScan() {
+ assertEmptyProjectBlocksAreSkipped(_segment, "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol LIMIT 25",
+ 1, 2, 1, 0);
+ }
+
+ /// Tail mode: the run scan reaches the project operator through {@code nextRow()}, which already tolerated empty
+ /// blocks. Pins that the two scan paths agree.
+ @Test
+ public void testEmptyProjectBlockDoesNotTruncateRunScan() {
+ assertEmptyProjectBlocksAreSkipped(_dupSegment,
+ "SELECT sortedCol, tailCol FROM testTable ORDER BY sortedCol, tailCol LIMIT 25", 1, 2, 1, 0);
+ }
+
+ /// The skip is a loop, not a single lookahead: several empty blocks in a row must all be skipped.
+ @Test
+ public void testConsecutiveEmptyProjectBlocksAreAllSkipped() {
+ assertEmptyProjectBlocksAreSkipped(_segment, "SELECT sortedCol FROM testTable ORDER BY sortedCol LIMIT 25", 1, 2,
+ 3, 0);
+ }
+
+ /// Empty blocks arriving immediately before exhaustion exercise the loop's other exit: the skip must fall through
+ /// to the project operator's {@code null} rather than spin or emit a phantom row. The limit exceeds the fixture so
+ /// the scan actually reaches the end instead of stopping on the row budget.
+ @Test
+ public void testEmptyProjectBlocksBeforeExhaustionEndTheScan() {
+ assertEmptyProjectBlocksAreSkipped(_segment, "SELECT sortedCol FROM testTable ORDER BY sortedCol LIMIT 100", 1, 0,
+ 0, 2);
+ }
+
+ /// The skip loop sits directly on top of the {@code limit + offset} budget bookkeeping, so cover a non-zero offset.
+ @Test
+ public void testEmptyProjectBlockIsSkippedWithOffset() {
+ assertEmptyProjectBlocksAreSkipped(_segment,
+ "SELECT sortedCol FROM testTable ORDER BY sortedCol LIMIT 15 OFFSET 10", 1, 2, 1, 0);
+ }
+
+ /// {@code nextRow()} rebuilds the phase-1 null bitmaps in the same branch that pulls the next non-empty block, so
+ /// the skip must not desynchronise a bitmap from the block it was built against. Uses the fixture that carries a
+ /// null in the order-by tail column, with null handling on.
+ @Test
+ public void testEmptyProjectBlockIsSkippedWithNullHandling() {
+ assertEmptyProjectBlocksAreSkipped(_nullTailSegment,
+ "SELECT tailCol, sortedCol, valCol FROM testTable ORDER BY sortedCol, tailCol LIMIT 25", 1, true, 2, 1, 0);
+ }
+
+ @Test
+ public void testZeroMatchSegmentEmitsOneSchemaBlock() {
+ assertZeroMatchEmitsOneSchemaBlock(_segment, "SELECT sortedCol FROM testTable WHERE sortedCol < 0 ORDER BY "
+ + "sortedCol", "SELECT sortedCol FROM testTable ORDER BY sortedCol", 1, false);
+ }
+
+ @Test
+ public void testZeroMatchSegmentEmitsOneSchemaBlockTwoPhase() {
+ assertZeroMatchEmitsOneSchemaBlock(_segment, "SELECT sortedCol, valCol FROM testTable WHERE sortedCol < 0 ORDER BY "
+ + "sortedCol", "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol", 1, false);
+ }
+
+ @Test
+ public void testZeroMatchSegmentEmitsOneSchemaBlockWithTailToSort() {
+ assertZeroMatchEmitsOneSchemaBlock(_dupSegment, "SELECT sortedCol, tailCol FROM testTable WHERE sortedCol < 0 "
+ + "ORDER BY sortedCol, tailCol", "SELECT sortedCol, tailCol FROM testTable ORDER BY sortedCol, tailCol", 1,
+ false);
+ }
+
+ @Test
+ public void testZeroMatchSegmentEmitsOneSchemaBlockWithNullHandling() {
+ assertZeroMatchEmitsOneSchemaBlock(_nullTailSegment, "SELECT tailCol, sortedCol, valCol FROM testTable WHERE "
+ + "sortedCol < 0 ORDER BY sortedCol, tailCol",
+ "SELECT tailCol, sortedCol, valCol FROM testTable ORDER BY sortedCol, tailCol", 1, true);
+ }
+
+ @Test
+ public void testMatchingSegmentEmitsNoTrailingEmptyBlock() {
+ assertNoEmptyBlockAmongEmittedBlocks(_segment, "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol", 1,
+ false);
+ }
+
+ @Test
+ public void testMatchingSegmentEmitsNoTrailingEmptyBlockWithTailToSort() {
+ assertNoEmptyBlockAmongEmittedBlocks(_dupSegment,
+ "SELECT sortedCol, tailCol, valCol FROM testTable ORDER BY sortedCol, tailCol", 1, false);
+ }
+
+ @Test
+ public void testMatchingSegmentEmitsNoTrailingEmptyBlockWithNullHandling() {
+ assertNoEmptyBlockAmongEmittedBlocks(_nullTailSegment,
+ "SELECT tailCol, sortedCol, valCol FROM testTable ORDER BY sortedCol, tailCol", 1, true);
+ }
+
+ /// A DESC query that cannot iterate the forward index in reverse falls through to the materialized operator,
+ /// reusing the streaming attempt's project or building a wider one. Both must match the plan that never entered
+ /// the streaming branch.
+ @Test
+ public void testDescIncompatibleFallbackMatchesTheMaterializedPlan() {
+ // Streaming project covers everything, so it is reused.
+ assertDescIncompatibleFallback(
+ "SET allowReverseOrder=false; SELECT sortedCol FROM testTable ORDER BY sortedCol DESC", true);
+ // Streaming project covers the order-by prefix only, so the fallback builds its own.
+ assertDescIncompatibleFallback(
+ "SET allowReverseOrder=false; SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol DESC", true);
+ // Zero-match: the reused project is driven to exhaustion without yielding a row.
+ assertDescIncompatibleFallback(
+ "SET allowReverseOrder=false; SELECT sortedCol FROM testTable WHERE sortedCol < 0 ORDER BY sortedCol DESC",
+ false);
+ // Not covered: allowReverseOrder on but withOrder(DESC) throwing. No fixture reaches it, and it lands on this
+ // same fall-through anyway.
+ }
+
+ private void assertDescIncompatibleFallback(@Language("sql") String query, boolean expectRows) {
+ SelectionResultsBlock block =
+ planWithCountedDataSourceLookups(query, SortedSelectionMergeMode.ON, new AtomicInteger()).nextBlock();
+ assertNotNull(block, "The fallback operator must emit a block");
+ List<Object[]> rows = block.getRows();
+ if (expectRows) {
+ assertFalse(rows.isEmpty(), "The fallback must still produce rows, else the assertions below are vacuous");
+ } else {
+ assertTrue(rows.isEmpty(), "A zero-match query must not produce rows");
+ }
+
+ SelectionResultsBlock baselineBlock =
+ planWithCountedDataSourceLookups(query, SortedSelectionMergeMode.OFF, new AtomicInteger()).nextBlock();
+ assertNotNull(baselineBlock);
+ List<Object[]> expectedRows = baselineBlock.getRows();
+ assertEquals(block.getDataSchema(), baselineBlock.getDataSchema(), "Schema mismatch for query: " + query);
+ assertEquals(rows.size(), expectedRows.size(), "Row count mismatch for query: " + query);
+ for (int i = 0; i < expectedRows.size(); i++) {
+ assertEquals(rows.get(i), expectedRows.get(i), "Row " + i + " mismatch for query: " + query);
+ }
+ }
+
+ /// Each {@link ProjectPlanNode} build looks up one data source per projected column, so a plan that builds the
+ /// project twice looks up more than the materialized plan, which builds it once.
+ @Test
+ public void testDescIncompatibleFallbackBuildsTheProjectOnce() {
+ @Language("sql") String query =
+ "SET allowReverseOrder=false; SELECT sortedCol FROM testTable ORDER BY sortedCol DESC";
+ AtomicInteger numStreamingLookups = new AtomicInteger();
+ AtomicInteger numMaterializedLookups = new AtomicInteger();
+ planWithCountedDataSourceLookups(query, SortedSelectionMergeMode.ON, numStreamingLookups);
+ planWithCountedDataSourceLookups(query, SortedSelectionMergeMode.OFF, numMaterializedLookups);
+ assertTrue(numMaterializedLookups.get() > 0, "Baseline must look up at least one data source");
+ assertEquals(numStreamingLookups.get(), numMaterializedLookups.get(),
+ "The DESC-incompatible fall-through built the project twice");
+ }
+
+ /// Plans {@code query} in {@code mode} against a segment counting data source lookups into {@code numLookups},
+ /// asserting it lands on the materialized DESC operator.
+ private Operator<SelectionResultsBlock> planWithCountedDataSourceLookups(@Language("sql") String query,
+ SortedSelectionMergeMode mode, AtomicInteger numLookups) {
+ IndexSegment countingSegment = mock(IndexSegment.class, delegatesTo(_segment));
+ // doAnswer, not when(...): building a when() on a delegating mock forwards to the real segment with nulls.
+ doAnswer(invocation -> {
+ numLookups.incrementAndGet();
+ return _segment.getDataSource(invocation.getArgument(0), invocation.getArgument(1));
+ }).when(countingSegment).getDataSource(anyString(), any());
+ QueryContext queryContext = QueryContextConverterUtils.getQueryContext(query);
+ queryContext.setSortedSelectionMergeMode(mode);
+ Operator<SelectionResultsBlock> operator =
+ new SelectionPlanNode(new SegmentContext(countingSegment), queryContext).run();
+ assertTrue(operator instanceof SelectionPartiallyOrderedByDescOperation,
+ "DESC without reverse order must fall back to the materialized DESC operator, got: "
+ + operator.getClass().getSimpleName());
+ return operator;
+ }
+
+ /// Every block the tail-to-sort path emits carries a structurally mutable row list. Defensive rather than a fix for
+ /// a reachable failure: no consumer of this operator adds to the list today, but {@link SelectionResultsBlock} is a
+ /// shared type and {@code SelectionOperatorUtils.mergeWithoutOrdering()} adds to the row list of the block it merges
+ /// into, which a fixed-size {@code Arrays.asList} view would reject.
+ @Test
+ public void testRunPathEmitsGrowableRowLists() {
+ QueryContext queryContext = QueryContextConverterUtils.getQueryContext(
+ "SELECT sortedCol, tailCol FROM testTable ORDER BY sortedCol, tailCol");
+ queryContext.setSortedSelectionMergeMode(SortedSelectionMergeMode.ON);
+ Operator<SelectionResultsBlock> operator =
+ new SelectionPlanNode(new SegmentContext(_dupSegment), queryContext).run();
+ assertTrue(operator instanceof StreamingSelectionOrderByOperator,
+ "Expected the streaming operator, got: " + operator.getClass().getSimpleName());
+
+ int numBlocks = 0;
+ SelectionResultsBlock block;
+ while ((block = operator.nextBlock()) != null) {
+ numBlocks++;
+ List<Object[]> rows = block.getRows();
+ assertFalse(rows.isEmpty(), "Block " + numBlocks + " must carry rows, else the assertion below is vacuous");
+ int numRows = rows.size();
+ rows.add(new Object[]{0, 0});
+ assertEquals(rows.size(), numRows + 1, "Block " + numBlocks + " did not accept an appended row");
+ }
+ // One block per run, so more than one block proves the tail-to-sort path (and therefore drainAscending) ran; the
+ // no-tail path emits a single block for this fixture.
+ assertTrue(numBlocks > 1, "Fixture must exercise the tail-to-sort run path, got " + numBlocks + " block(s)");
+ }
+
+ /// The counterpart to the zero-match tests: a segment that does match rows must never emit an empty block, so the
+ /// guarantee reads "at least one block" and not "always an extra one".
+ private void assertNoEmptyBlockAmongEmittedBlocks(IndexSegment segment, @Language("sql") String query,
+ int numSortedExpressions, boolean nullHandling) {
+ StreamingSelectionOrderByOperator operator =
+ buildOperatorWithInjectableProject(segment, query, numSortedExpressions, nullHandling, 0, 0, 0);
+ int numBlocks = 0;
+ SelectionResultsBlock block;
+ while ((block = operator.nextBlock()) != null) {
+ numBlocks++;
+ assertFalse(block.getRows().isEmpty(),
+ "Block " + numBlocks + " is empty: the schema block must only be emitted when no rows were emitted at all");
+ }
+ assertTrue(numBlocks > 1, "Fixture must span more than one block, otherwise a trailing block cannot be observed");
+ }
+
+ /// A segment matching no rows must still emit exactly one block before signalling exhaustion: empty, but carrying
+ /// the same {@link DataSchema} the identical query produces when it does match rows. Consumers therefore never have
+ /// to reconstruct a schema the segment already knows.
+ private void assertZeroMatchEmitsOneSchemaBlock(IndexSegment segment, @Language("sql") String zeroMatchQuery,
+ @Language("sql") String matchingQuery, int numSortedExpressions, boolean nullHandling) {
+ SelectionResultsBlock matchingBlock =
+ buildOperatorWithInjectableProject(segment, matchingQuery, numSortedExpressions, nullHandling, 0, 0, 0)
+ .nextBlock();
+ assertNotNull(matchingBlock);
+ assertFalse(matchingBlock.getRows().isEmpty(), "The control query must actually match rows");
+ DataSchema expectedSchema = matchingBlock.getDataSchema();
+ assertNotNull(expectedSchema);
+
+ StreamingSelectionOrderByOperator operator =
+ buildOperatorWithInjectableProject(segment, zeroMatchQuery, numSortedExpressions, nullHandling, 0, 0, 0);
+ SelectionResultsBlock block = operator.nextBlock();
+ assertNotNull(block, "A zero-match segment must still emit one block carrying the schema");
+ assertTrue(block.getRows().isEmpty(), "A zero-match segment must not emit rows");
+ assertEquals(block.getDataSchema(), expectedSchema,
+ "The zero-match schema must equal the schema the same query produces when it matches rows");
+ assertNull(operator.nextBlock(), "Exactly one block may precede exhaustion");
+ }
+
+ /// Drives the operator twice over {@code segment}: once against the real project operator, once against one that
+ /// splices zero-document blocks into its output, and asserts the emitted rows are identical.
+ ///
+ /// <p>No doc-id-set operator produces an empty block today, so the case has to be injected. The project operator is
+ /// built with a deliberately small {@code maxDocsPerCall} so the scan spans several blocks.
+ ///
+ /// @param injectBeforeRealBlock 1-based index of the real block to splice empties in front of; 0 for none
+ /// @param numEmptyBlocks how many consecutive empty blocks to splice in at that point
+ /// @param numTrailingEmptyBlocks how many empty blocks to emit after the last real block, before exhaustion
+ private void assertEmptyProjectBlocksAreSkipped(IndexSegment segment, @Language("sql") String query,
+ int numSortedExpressions, int injectBeforeRealBlock, int numEmptyBlocks, int numTrailingEmptyBlocks) {
+ assertEmptyProjectBlocksAreSkipped(segment, query, numSortedExpressions, false, injectBeforeRealBlock,
+ numEmptyBlocks, numTrailingEmptyBlocks);
+ }
+
+ /// As above, with explicit control over null handling.
+ private void assertEmptyProjectBlocksAreSkipped(IndexSegment segment, @Language("sql") String query,
+ int numSortedExpressions, boolean nullHandling, int injectBeforeRealBlock, int numEmptyBlocks,
+ int numTrailingEmptyBlocks) {
+ List<Object[]> expected =
+ drain(buildOperatorWithInjectableProject(segment, query, numSortedExpressions, nullHandling, 0, 0, 0));
+ assertTrue(expected.size() > MAX_DOCS_PER_PROJECT_BLOCK,
+ "Fixture must span more than one project block, otherwise the injection is not exercised");
+ List<Object[]> actual = drain(
+ buildOperatorWithInjectableProject(segment, query, numSortedExpressions, nullHandling, injectBeforeRealBlock,
+ numEmptyBlocks, numTrailingEmptyBlocks));
+ assertEquals(actual.size(), expected.size(),
+ "Row count changed when empty project blocks were injected, so the scan was truncated");
+ for (int i = 0; i < expected.size(); i++) {
+ assertEquals(actual.get(i), expected.get(i), "Row " + i + " mismatch for query: " + query);
+ }
+ }
+
+ /// Builds a {@link StreamingSelectionOrderByOperator} directly (rather than through {@link SelectionPlanNode}) so
+ /// the project operator can be wrapped. Passing 0 for every injection parameter yields the undecorated operator.
+ private StreamingSelectionOrderByOperator buildOperatorWithInjectableProject(IndexSegment segment,
+ @Language("sql") String query, int numSortedExpressions, boolean nullHandling, int injectBeforeRealBlock,
+ int numEmptyBlocks, int numTrailingEmptyBlocks) {
+ QueryContext queryContext = QueryContextConverterUtils.getQueryContext(query);
+ queryContext.setNullHandlingEnabled(nullHandling);
+ queryContext.setSortedSelectionMergeMode(SortedSelectionMergeMode.ON);
+ List<ExpressionContext> expressions = SelectionOperatorUtils.extractExpressions(queryContext, segment);
+ List<OrderByExpressionContext> orderByExpressions = queryContext.getOrderByExpressions();
+ assertNotNull(orderByExpressions);
+ // Mirrors SelectionPlanNode: only the order-by expressions are scanned when there is a second phase.
+ List<ExpressionContext> projectExpressions = expressions;
+ if (expressions.size() > orderByExpressions.size()) {
+ projectExpressions = new ArrayList<>(orderByExpressions.size());
+ for (OrderByExpressionContext orderByExpression : orderByExpressions) {
+ projectExpressions.add(orderByExpression.getExpression());
+ }
+ }
+ BaseProjectOperator<?> projectOperator =
+ new ProjectPlanNode(new SegmentContext(segment), queryContext, projectExpressions,
+ MAX_DOCS_PER_PROJECT_BLOCK).run();
+ BaseProjectOperator<?> effective =
+ injectBeforeRealBlock == 0 && numTrailingEmptyBlocks == 0
+ ? projectOperator
+ : new EmptyBlockInjectingProjectOperator(projectOperator, injectBeforeRealBlock, numEmptyBlocks,
+ numTrailingEmptyBlocks);
+ return new StreamingSelectionOrderByOperator(segment, queryContext, expressions, effective, numSortedExpressions);
+ }
+
+ private static List<Object[]> drain(StreamingSelectionOrderByOperator operator) {
+ List<Object[]> rows = new ArrayList<>();
+ SelectionResultsBlock block;
+ while ((block = operator.nextBlock()) != null) {
+ rows.addAll(block.getRows());
+ }
+ return rows;
+ }
+
+ /// Forwards everything to a real project operator, but splices zero-document blocks into its output:
+ /// {@code numEmptyBlocks} of them just before the {@code injectBeforeRealBlock}-th real block, and
+ /// {@code numTrailingEmptyBlocks} after the last real block but before exhaustion.
+ ///
+ /// <p>Every real block is still delivered, merely deferred, so the decorator provably drops nothing of its own -
+ /// any row loss observed by a test is the operator under test truncating its scan.
+ private static class EmptyBlockInjectingProjectOperator extends BaseProjectOperator<ValueBlock> {
+ private final BaseProjectOperator<?> _delegate;
+ private final int _injectBeforeRealBlock;
+ private final int _numEmptyBlocks;
+ private final int _numTrailingEmptyBlocks;
+ private final Deque<ValueBlock> _pending = new ArrayDeque<>();
+ private int _realBlocksSeen;
+ private ValueBlock _lastRealBlock;
+ private boolean _injected;
+ private boolean _trailingEmitted;
+
+ EmptyBlockInjectingProjectOperator(BaseProjectOperator<?> delegate, int injectBeforeRealBlock, int numEmptyBlocks,
+ int numTrailingEmptyBlocks) {
+ _delegate = delegate;
+ _injectBeforeRealBlock = injectBeforeRealBlock;
+ _numEmptyBlocks = numEmptyBlocks;
+ _numTrailingEmptyBlocks = numTrailingEmptyBlocks;
+ }
+
+ @Override
+ protected ValueBlock getNextBlock() {
+ if (!_pending.isEmpty()) {
+ return _pending.poll();
+ }
+ ValueBlock block = _delegate.nextBlock();
+ if (block == null) {
+ // The empty blocks shadow the last real block so their value sets stay usable; there is nothing to shadow if
+ // the delegate never produced one, in which case the stream simply ends.
+ if (_trailingEmitted || _numTrailingEmptyBlocks == 0 || _lastRealBlock == null) {
+ return null;
+ }
+ _trailingEmitted = true;
+ for (int i = 0; i < _numTrailingEmptyBlocks; i++) {
+ _pending.add(new EmptyValueBlock(_lastRealBlock));
+ }
+ return _pending.poll();
+ }
+ _lastRealBlock = block;
+ _realBlocksSeen++;
+ if (!_injected && _realBlocksSeen == _injectBeforeRealBlock) {
+ _injected = true;
+ for (int i = 0; i < _numEmptyBlocks; i++) {
+ _pending.add(new EmptyValueBlock(block));
+ }
+ _pending.add(block);
+ return _pending.poll();
+ }
+ return block;
+ }
+
+ @Override
+ public Map<String, ColumnContext> getSourceColumnContextMap() {
+ return _delegate.getSourceColumnContextMap();
+ }
+
+ @Override
+ public ColumnContext getResultColumnContext(ExpressionContext expression) {
+ return _delegate.getResultColumnContext(expression);
+ }
+
+ @Override
+ public BaseProjectOperator<ValueBlock> withOrder(DocIdOrder newOrder) {
+ throw new UnsupportedOperationException("Test decorator does not support reordering");
+ }
+
+ @Override
+ public boolean isCompatibleWith(DocIdOrder order) {
+ return _delegate.isCompatibleWith(order);
+ }
+
+ @Override
+ public ExecutionStatistics getExecutionStatistics() {
+ return _delegate.getExecutionStatistics();
+ }
+
+ @Override
+ public List<? extends Operator> getChildOperators() {
+ return List.of(_delegate);
+ }
+
+ @Override
+ public String toExplainString() {
+ return "EMPTY_BLOCK_INJECTING_PROJECT";
+ }
+ }
+
+ /// A block reporting zero documents. Value sets are forwarded to the real block it shadows so that a consumer which
+ /// builds its fetchers before checking the document count does not fail for the wrong reason.
+ private static class EmptyValueBlock implements ValueBlock {
+ private final ValueBlock _delegate;
+
+ EmptyValueBlock(ValueBlock delegate) {
+ _delegate = delegate;
+ }
+
+ @Override
+ public int getNumDocs() {
+ return 0;
+ }
+
+ @Override
+ public int[] getDocIds() {
+ return new int[0];
+ }
+
+ @Override
+ public BlockValSet getBlockValueSet(ExpressionContext expression) {
+ return _delegate.getBlockValueSet(expression);
+ }
+
+ @Override
+ public BlockValSet getBlockValueSet(String column) {
+ return _delegate.getBlockValueSet(column);
+ }
+
+ @Override
+ public BlockValSet getBlockValueSet(String[] paths) {
+ return _delegate.getBlockValueSet(paths);
+ }
+ }
+
/// Runs {@code query} twice over {@code segment} - once with the streaming hint on, once off - and asserts the
/// concatenated streaming blocks equal the materialized operator's single block, cell by cell. Returns the (now
/// exhausted) streaming operator so callers can make extra assertions on its execution statistics.
@@ -338,7 +803,7 @@ public class StreamingSelectionOrderByOperatorTest {
// Streaming path.
QueryContext streamingContext = QueryContextConverterUtils.getQueryContext(query);
streamingContext.setNullHandlingEnabled(nullHandling);
- streamingContext.setSortedSelectionMergeEnabled(true);
+ streamingContext.setSortedSelectionMergeMode(SortedSelectionMergeMode.ON);
Operator<SelectionResultsBlock> streamingOperator =
new SelectionPlanNode(new SegmentContext(segment), streamingContext).run();
assertTrue(streamingOperator instanceof StreamingSelectionOrderByOperator,
@@ -382,7 +847,7 @@ public class StreamingSelectionOrderByOperatorTest {
boolean nullHandling) {
QueryContext queryContext = QueryContextConverterUtils.getQueryContext(query);
queryContext.setNullHandlingEnabled(nullHandling);
- queryContext.setSortedSelectionMergeEnabled(true);
+ queryContext.setSortedSelectionMergeMode(SortedSelectionMergeMode.ON);
Operator<SelectionResultsBlock> operator = new SelectionPlanNode(new SegmentContext(segment), queryContext).run();
List<Object[]> rows = new ArrayList<>();
SelectionResultsBlock block;
diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
index 0af0e797710..4a8d90207b1 100644
--- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
+++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
@@ -930,8 +930,15 @@ public class CommonConstants {
/// there is no such stage and the client would observe duplicate rows across flush windows.
public static final String STREAMING_DISTINCT_FLUSH_THRESHOLD = "streamingDistinctFlushThreshold";
- /// Opt-in: use the streaming k-way-merge selection ORDER BY combine over sorted segments.
- public static final String SORTED_SELECTION_MERGE_ENABLED = "sortedSelectionMergeEnabled";
+ /// Selects the streaming k-way-merge selection ORDER BY combine over sorted segments. One of
+ /// [Server.SortedSelectionMergeMode] (case-insensitive); defaults to
+ /// [Server#DEFAULT_SORTED_SELECTION_MERGE_MODE].
+ public static final String SORTED_SELECTION_MERGE_MODE = "sortedSelectionMergeMode";
+ /// Minimum fraction of segments that must be physically sorted on the leading ORDER BY column for
+ /// [Server.SortedSelectionMergeMode#AUTO] to select the streaming merge. Overrides
+ /// [Server#DEFAULT_SORTED_SELECTION_MERGE_AUTO_MIN_SORTED_RATIO].
+ public static final String SORTED_SELECTION_MERGE_AUTO_MIN_SORTED_RATIO =
+ "sortedSelectionMergeAutoMinSortedRatio";
/// Output block size (rows) for the streaming selection ORDER BY combine.
public static final String SORTED_SELECTION_MERGE_BLOCK_SIZE = "sortedSelectionMergeBlockSize";
@@ -1611,6 +1618,38 @@ public class CommonConstants {
// Sequential combine utilizes N + 1 threads, N for processing and 1 for merge. Pair-wise only uses N.
public static final int DEFAULT_SORT_AGGREGATE_SEQUENTIAL_COMBINE_NUM_SEGMENTS_THRESHOLD =
Runtime.getRuntime().availableProcessors();
+
+ /// Selects the streaming k-way-merge selection ORDER BY combine
+ /// ([Broker.Request.QueryOptionKey#SORTED_SELECTION_MERGE_MODE]).
+ ///
+ /// The merge runs single-threaded on the consumer thread, so it trades `maxExecutionThreads` worth of parallel
+ /// segment processing for the ability to emit sorted blocks incrementally. That trade only pays off when the
+ /// segments are physically sorted on the leading ORDER BY column, since an unsorted child must do a full scan
+ /// plus top-K synchronously before the merge can take its first row.
+ public enum SortedSelectionMergeMode {
+ /// Never use the streaming merge. Selection ORDER BY keeps its existing combine operators.
+ OFF,
+ /// Always use the streaming merge, including when the leading ORDER BY expression is not an identifier and no
+ /// segment is sorted. Children then materialize their top-K eagerly and the merge is strictly slower than
+ /// [#OFF]; intended for benchmarking and A/B, not for production traffic.
+ ON,
+ /// Use the streaming merge only when at least
+ /// [Broker.Request.QueryOptionKey#SORTED_SELECTION_MERGE_AUTO_MIN_SORTED_RATIO] of the queried segments are
+ /// physically sorted on the leading ORDER BY column. Decided from segment metadata alone, without acquiring
+ /// any segment.
+ AUTO
+ }
+
+ public static final SortedSelectionMergeMode DEFAULT_SORTED_SELECTION_MERGE_MODE =
+ SortedSelectionMergeMode.OFF;
+
+ /// Default for [Broker.Request.QueryOptionKey#SORTED_SELECTION_MERGE_AUTO_MIN_SORTED_RATIO].
+ ///
+ /// Consuming segments are never reported as sorted ([org.apache.pinot.segment.spi.datasource.DataSourceMetadata]
+ /// on a mutable segment hardcodes `isSorted()` to false), so a realtime table always has a few unsorted segments
+ /// among many sorted ones. This default is chosen to tolerate that while still rejecting a mostly-unsorted table.
+ /// It is a starting point rather than a measured optimum.
+ public static final double DEFAULT_SORTED_SELECTION_MERGE_AUTO_MIN_SORTED_RATIO = 0.8;
public static final String CONFIG_OF_MSE_MIN_GROUP_TRIM_SIZE = MSE_CONFIG_PREFIX + ".min.group.trim.size";
// Match the value of GroupByUtils.DEFAULT_MIN_NUM_GROUPS
public static final int DEFAULT_MSE_MIN_GROUP_TRIM_SIZE = 5000;
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
(pinot) 04/07: Address second-round review on the streaming selection ORDER BY combine
Posted by xi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
xiangfu0 pushed a commit to branch xiangfu0/codex1/sorted-exchange-leaf
in repository https://gitbox.apache.org/repos/asf/pinot.git
commit a8065e5459d30ee62f9d7b2ad091f7aee580a622
Author: rohity <ro...@uber.com>
AuthorDate: Sun Sep 27 14:07:54 2026 +0000
Address second-round review on the streaming selection ORDER BY combine
Test segment release and schema-mismatch handling. Nothing exercised either
path before: the existing tests plan the combine from plain segment plan
nodes, so the acquire/release calls were no-ops and a leaked segment would
have passed, and every segment shared one schema, so the mismatch branch
never ran. The new tests wrap each real streaming leaf in a counting
AcquireReleaseColumnsSegmentOperator, the way the prefetch path plans it,
and assert every acquired segment is released exactly once when the merge
drains, when the limit stops it with cursors still open, on stop()
mid-merge (including a repeated stop()), and when a child throws an
exception or an Error. A mismatched schema on one block must drop the rest
of that segment, leave every other segment whole, and report one
deduplicated MERGE_RESPONSE error. Each was checked against an operator
with the corresponding release or drop removed. Document why the cursor
drops the rest of the segment on a mismatch rather than one block: it is
the same unit the merger drops, since there a block is a segment's whole
result.
Add server config defaults for the merge options.
sortedSelectionMergeBlockSize and sortedSelectionMergeAutoMinSortedRatio
could only be changed per query. Add
pinot.server.query.executor.sorted.selection.merge.block.size and
pinot.server.query.executor.sorted.selection.merge.auto.min.sorted.ratio,
read and validated in InstancePlanMakerImplV2#init and applied whenever the
query does not set the option, as numGroupsLimit already is. The block size
default moves from Broker to Server, since only the server reads it.
applyQueryOptions now always writes both values, so the AUTO threshold
tests set the ratio through the query option instead of the query context
setter. Also bound CombineSlowOperatorsTest's deadline test with a timeout:
a regression there drives a SlowOperator that sleeps for an hour, which
would hang the build instead of failing it.
Expose the combine's settings in the explain plan. It now reports block
size, frontier pruning, tie deferral and the segment / sorted-segment
counts as explain attributes, so an MSE explain with explainAskingServers
shows which path ran and whether pruning was active. Tests pin that
numSegmentsMatched excludes never-activated segments.
Use markdown syntax in the /// javadoc added by this change: [Foo],
backticks, **bold** and blank-line paragraphs instead of {@link}, {@code},
<b> and <p>, matching the rest of the module.
---
.../StreamingSelectionOrderByCombineOperator.java | 129 ++++---
.../query/StreamingSelectionOrderByOperator.java | 48 +--
.../core/plan/maker/InstancePlanMakerImplV2.java | 32 +-
.../core/query/request/context/QueryContext.java | 3 +-
.../operator/combine/CombineSlowOperatorsTest.java | 9 +-
...reamingSelectionOrderByCombineOperatorTest.java | 381 +++++++++++++++++++--
.../StreamingSelectionOrderByOperatorTest.java | 68 ++--
.../plan/maker/InstancePlanMakerImplV2Test.java | 67 +++-
.../apache/pinot/spi/utils/CommonConstants.java | 19 +-
9 files changed, 601 insertions(+), 155 deletions(-)
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperator.java
index 75560d3715b..5597a06bcfc 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperator.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperator.java
@@ -32,6 +32,7 @@ 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;
@@ -49,17 +50,17 @@ 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
+/// 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
-/// {@link StreamingSelectionOrderByOperator}, which is pulled lazily one run/block at a time.
+/// [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
+/// [SelectionResultsBlock]-producing operator such as `SelectionOrderByOperator`); the cursor reads that
/// one block and iterates its rows.
///
-/// A {@link PriorityQueue} of {@link SegmentCursor} ordered by the {@link OrderByComparatorFactory} comparator on
+/// 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
@@ -67,39 +68,43 @@ import org.slf4j.LoggerFactory;
/// 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(SegmentCursor)} for the
+/// (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.
/// 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.
+/// 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
-/// {@link AcquireReleaseColumnsSegmentOperator} on activation and releases it only when its child operator is fully
+/// [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()}
+/// [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 {@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()}.
+/// 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
-/// {@link org.apache.pinot.core.query.executor.ResultsBlockStreamer} is present (the
-/// MSE leaf, driven by {@link org.apache.pinot.core.operator.streaming.StreamingInstanceResponseOperator}, and the
-/// gRPC streaming endpoint). The merge emits many bounded {@link SelectionResultsBlock}s from successive
-/// {@link #getNextBlock()} calls followed by a final {@link MetadataResultsBlock}. On the blocking single-stage
-/// path the caller expects one complete block from a single {@link #getNextBlock()} call, so nothing here could
-/// stream: the merge would accumulate the whole {@code limit + offset} result while holding every activated cursor's
-/// materialized block alive. That path keeps {@link MinMaxValueBasedSelectionOrderByCombineOperator}.
+/// [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 {@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. The base entry points that it replaces - {@link #processSegments()} and
-/// {@link #isQuerySatisfied(SelectionResultsBlock, Object)} - are overridden to fail loud. The base
-/// {@code Phaser} (which exists only to fence worker threads against segment release) is intentionally bypassed because
+/// **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.
@@ -111,10 +116,12 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
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 {@link #sortsBeyond}.
+ /// 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;
@@ -168,6 +175,7 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
// 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) {
@@ -178,7 +186,11 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
operator.getIndexSegment().getDataSource(firstOrderByColumn, queryContext.getSchema())
.getDataSourceMetadata();
_sortedCursors[i] = new SegmentCursor(operator, metadata.getMinValue(), metadata.getMaxValue());
+ if (metadata.isSorted()) {
+ numSortedSegments++;
+ }
}
+ _numSortedSegments = numSortedSegments;
if (firstOrderByColumn != null) {
sortCursorsByMinMax();
}
@@ -190,7 +202,7 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
/// 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}).
+ /// always be processed (mirrors [MinMaxValueBasedSelectionOrderByCombineOperator]).
private void sortCursorsByMinMax() {
if (_asc) {
Arrays.sort(_sortedCursors, (o1, o2) -> {
@@ -220,7 +232,23 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
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
+ /// 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() {
@@ -234,7 +262,7 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
releaseAllCursors();
}
- /// The base worker-thread entry point must never run here ({@link #start()} is a no-op). Fail loud if it ever does.
+ /// 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(
@@ -329,12 +357,12 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
/// Activates not-yet-active cursors whose min/max bound can still reach the merge frontier.
///
- /// Cursors are visited in min/max order and {@code _nextToActivate} advances only on a real activation, so a
- /// {@code break} defers the current cursor to a later call with a risen frontier rather than skipping it. Deferring
+ /// Cursors are visited in min/max order and `_nextToActivate` advances only on a real activation, so a
+ /// `break` defers the current cursor to a later call with a risen frontier rather than skipping it. Deferring
/// is safe because the first order-by column is the primary sort key: a cursor whose range starts past the frontier
/// holds no row sorting before it. With no frontier known yet, activation is forced. Under a single order-by
/// expression a cursor whose bound merely *ties* the frontier is deferred too, which is what keeps a long run of
- /// equal segment minima from activating every segment at once; see {@link #sortsBeyond}.
+ /// equal segment minima from activating every segment at once; see [#sortsBeyond].
///
/// The frontier is the row about to be emitted. Understating it defers a cursor that could have supplied that row
/// and the merge emits out of order, where overstating it only activates a segment early. The leader is retained
@@ -366,11 +394,11 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
/// be compared, so this reports `false` and the caller activates. Tests column 0 only, as the pruning bound always
/// has, rather than the full-row comparator -- which would put every order-by column on the per-row path.
///
- /// Under a single order-by expression ({@code _deferTiedCursors}) a *tie* also defers, which only postpones a
- /// cursor: {@code _nextToActivate} does not advance, and the caller force-activates once heap and leader drain.
- /// Since {@code bound} bounds every row in the segment, a deferred cursor holds no row sorting strictly before an
+ /// Under a single order-by expression (`_deferTiedCursors`) a *tie* also defers, which only postpones a
+ /// cursor: `_nextToActivate` does not advance, and the caller force-activates once heap and leader drain.
+ /// Since `bound` bounds every row in the segment, a deferred cursor holds no row sorting strictly before an
/// emitted one -- only ties, interchangeable when column 0 is the whole sort key. Note this is *not*
- /// {@link MinMaxValueBasedSelectionOrderByCombineOperator}'s justification: its bound is a complete top-K's k-th
+ /// [MinMaxValueBasedSelectionOrderByCombineOperator]'s justification: its bound is a complete top-K's k-th
/// row, this one is the live frontier. It changes which tied rows are returned, already documented as arbitrary.
private boolean sortsBeyond(Comparable bound, @Nullable SegmentCursor leader, @Nullable SegmentCursor top) {
Object leaderHead = leader != null ? leader.currentHead()[0] : null;
@@ -404,7 +432,7 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
}
}
- /// Returns the accumulated output rows as a sorted {@link SelectionResultsBlock} and resets the output buffer. The
+ /// Returns the accumulated output rows as a sorted [SelectionResultsBlock] and resets the output buffer. The
/// block carries the comparator so the broker-side n-way reduce stays correct. Execution stats are not attached
/// here: they go on the terminal metadata block.
private BaseResultsBlock flushDataBlock() {
@@ -453,9 +481,9 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
}
/// Iterates a single segment's locally-sorted rows, pulling blocks lazily from its operator. Streaming-backed cursors
- /// loop until the operator returns {@code null}; single-block-backed cursors read exactly one block. The segment is
+ /// loop until the operator returns `null`; single-block-backed cursors read exactly one block. The segment is
/// acquired on activation and released once exhausted or when the combine finishes (see
- /// {@link StreamingSelectionOrderByCombineOperator}).
+ /// [StreamingSelectionOrderByCombineOperator]).
private class SegmentCursor {
private final Operator<BaseResultsBlock> _operator;
@Nullable
@@ -478,14 +506,14 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
_maxValue = maxValue;
}
- /// Returns the current head row to be merged next, or {@code null} if not activated or exhausted.
+ /// Returns the current head row to be merged next, or `null` if not activated or exhausted.
@Nullable
Object[] currentHead() {
return _head;
}
/// Acquires the segment, resolves whether the child is the lazy streaming operator, and reads the first block.
- /// After this call {@link #currentHead()} returns the first row, or {@code null} if the segment contributes
+ /// After this call [#currentHead()] returns the first row, or `null` if the segment contributes
/// nothing.
void activate() {
acquireSegment();
@@ -495,7 +523,7 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
}
/// Advances past the current head, pulling the next block lazily for streaming cursors. Releases the segment
- /// when the cursor is exhausted; afterwards {@link #currentHead()} returns {@code null}.
+ /// when the cursor is exhausted; afterwards [#currentHead()] returns `null`.
void advance() {
_pos++;
if (_pos < _rows.size()) {
@@ -510,14 +538,15 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
}
/// Loads the next non-empty block of rows from the operator, capturing the combine-level data schema on first
- /// sight. Returns {@code false} when the operator is exhausted (no more rows). Streaming-backed operators emit
- /// one run/block per call and {@code null} when done; single-block operators emit a single block and must not be
+ /// sight. Returns `false` when the operator is exhausted (no more rows). Streaming-backed operators emit
+ /// one run/block per call and `null` when done; single-block operators emit a single block and must not be
/// called again afterwards, so an empty/null block from a single-block child is treated as exhausted.
///
/// A block whose schema differs from the one already captured is dropped and reported, mirroring
- /// {@link org.apache.pinot.core.operator.combine.merger.SelectionOrderByResultsBlockMerger}. Segments on a server
+ /// [org.apache.pinot.core.operator.combine.merger.SelectionOrderByResultsBlockMerger]. Segments on a server
/// can disagree on schema mid-reload (a newly added column exists only in reloaded segments), and merging rows of
- /// differing width under one schema would corrupt the result rather than fail.
+ /// differing width under one schema would corrupt the result rather than fail. The cursor is exhausted rather than
+ /// skipping the one block, which drops the same unit the merger does: there, a block is a segment's whole result.
private boolean pullBlock() {
while (true) {
SelectionResultsBlock block = nextBlock();
@@ -562,8 +591,8 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
}
}
- /// Returns whether the underlying child operator is the lazy {@link StreamingSelectionOrderByOperator}. Must be
- /// called after the first {@link Operator#nextBlock()}, which is what materializes a wrapped child.
+ /// Returns whether the underlying child operator is the lazy [StreamingSelectionOrderByOperator]. Must be
+ /// called after the first [Operator#nextBlock()], which is what materializes a wrapped child.
private boolean isStreamingChild() {
Operator underlying = _operator;
if (_operator instanceof AcquireReleaseColumnsSegmentOperator) {
@@ -583,7 +612,7 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
_acquired = true;
}
- /// Releases the segment if still held. Idempotent: safe to call from {@link #exhaust()} and combine cleanup.
+ /// Releases the segment if still held. Idempotent: safe to call from [#exhaust()] and combine cleanup.
private void release() {
if (_acquired) {
if (_operator instanceof AcquireReleaseColumnsSegmentOperator) {
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperator.java
index 10dbd97ccde..58f5f37ef7b 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperator.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperator.java
@@ -61,32 +61,32 @@ 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
+/// 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 {@code projectOperator.isCompatibleWith(DocIdOrder.fromAsc(asc))}).
+/// the query order (the caller must guarantee `projectOperator.isCompatibleWith(DocIdOrder.fromAsc(asc))`).
///
-/// Every instance emits at least one block before {@code null}: a segment matching no rows emits a single empty block
-/// carrying the {@link DataSchema}, so consumers never have to reconstruct a schema the segment already knows.
+/// 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** ({@code numSortedExpressions == numOrderByExpressions}, e.g. {@code ORDER BY sorted}):
+/// - **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 {@code limit + offset} budget).
-/// - **Tail to sort** ({@code numSortedExpressions < numOrderByExpressions}, e.g.
-/// {@code ORDER BY sorted, other}):
+/// (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
-/// {@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).
+/// `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 {@link SelectionOrderByOperator} it preserves the two-phase projection optimization: when there are output
+/// 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 {@link #getNextBlock()} calls and is **not** thread-safe; a single consumer
+/// 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";
@@ -231,7 +231,7 @@ public class StreamingSelectionOrderByOperator extends BaseOperator<SelectionRes
}
/// 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, and otherwise a
+ /// 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
@@ -273,7 +273,7 @@ public class StreamingSelectionOrderByOperator extends BaseOperator<SelectionRes
}
/// 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
+ /// `limit + offset` rows by the full comparator, and return them sorted. Returns `null` when
/// exhausted.
@Nullable
private List<Object[]> nextRun() {
@@ -314,8 +314,8 @@ public class StreamingSelectionOrderByOperator extends BaseOperator<SelectionRes
}
/// Pulls the next row of the forward scan (across project blocks), materialized as an
- /// {@code Object[_numExpressions]}. For two-phase the document id is stashed at index
- /// {@code _numOrderByExpressions} (overwritten in the second pass). Returns {@code null} when the project operator
+ /// `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() {
@@ -350,7 +350,7 @@ public class StreamingSelectionOrderByOperator extends BaseOperator<SelectionRes
return materializeRow(_currentFetcher, _currentDocIds, _currentNullBitmaps, rowId);
}
- /// Pulls the next project block carrying documents, skipping any that carry none. Returns {@code null} only when
+ /// 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.
///
@@ -410,7 +410,7 @@ public class StreamingSelectionOrderByOperator extends BaseOperator<SelectionRes
/// 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 {@link SelectionOrderByOperator#computePartiallyOrdered()}.
+ /// instances, mirroring [SelectionOrderByOperator#computePartiallyOrdered()].
private void fetchNonOrderByColumns(List<Object[]> rows) {
int numRows = rows.size();
RoaringBitmap docIds = new RoaringBitmap();
@@ -478,9 +478,9 @@ public class StreamingSelectionOrderByOperator extends BaseOperator<SelectionRes
}
/// Resolves a non-order-by expression's result type without building the phase-2 pipeline, for the zero-match case
- /// where there is nothing to fetch. This reproduces {@link TransformOperator#getResultColumnContext} exactly -- that
+ /// where there is nothing to fetch. This reproduces [TransformOperator#getResultColumnContext] exactly -- that
/// method resolves against its project operator's source column contexts, which for phase 2 are
- /// {@link #_phase2DataSourceMap} -- so the schema is identical to the one the fetch path builds, at metadata cost
+ /// [#_phase2DataSourceMap] -- so the schema is identical to the one the fetch path builds, at metadata cost
/// only. Worth the indirection because a selective filter can leave many segments of a table matching nothing, and
/// each would otherwise stand up a full projection and transform operator to read types it already knows.
private ColumnContext resolveResultColumnContext(ExpressionContext expression) {
@@ -558,8 +558,8 @@ public class StreamingSelectionOrderByOperator extends BaseOperator<SelectionRes
numTotalDocs);
}
- /// Reports scan cost to the shared {@link QueryScanCostContext} so scan-based query killing
- /// ({@link org.apache.pinot.core.common.Operator#nextBlock()} -> {@code checkScanBasedKilling}) can see this
+ /// Reports scan cost to the shared [QueryScanCostContext] so scan-based query killing
+ /// ([org.apache.pinot.core.common.Operator#nextBlock()] -> `checkScanBasedKilling`) can see this
/// operator's work. Without this the killer is blind on the streaming path, unlike every other selection operator.
private void reportScanCost(int numDocsScanned, long numEntriesScannedPostFilter) {
QueryScanCostContext scanCost = getScanCostContext();
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java b/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java
index e0c8c541184..f00ffb4aa18 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java
@@ -116,6 +116,9 @@ public class InstancePlanMakerImplV2 implements PlanMaker {
private int _minSegmentGroupTrimSize = Server.DEFAULT_QUERY_EXECUTOR_MIN_SEGMENT_GROUP_TRIM_SIZE;
private int _minServerGroupTrimSize = Server.DEFAULT_QUERY_EXECUTOR_MIN_SERVER_GROUP_TRIM_SIZE;
private int _groupByTrimThreshold = Server.DEFAULT_QUERY_EXECUTOR_GROUPBY_TRIM_THRESHOLD;
+ // Server-wide defaults for the streaming selection ORDER BY merge; the query options override them
+ private double _sortedSelectionMergeAutoMinSortedRatio = Server.DEFAULT_SORTED_SELECTION_MERGE_AUTO_MIN_SORTED_RATIO;
+ private int _sortedSelectionMergeBlockSize = Server.DEFAULT_SORTED_SELECTION_MERGE_BLOCK_SIZE;
@Override
public void init(PinotConfiguration queryExecutorConfig) {
@@ -146,11 +149,24 @@ public class InstancePlanMakerImplV2 implements PlanMaker {
Server.DEFAULT_QUERY_EXECUTOR_GROUPBY_TRIM_THRESHOLD);
Preconditions.checkState(_groupByTrimThreshold > 0,
"Invalid configurable: groupByTrimThreshold: %d must be positive", _groupByTrimThreshold);
+ _sortedSelectionMergeAutoMinSortedRatio =
+ queryExecutorConfig.getProperty(Server.SORTED_SELECTION_MERGE_AUTO_MIN_SORTED_RATIO,
+ Server.DEFAULT_SORTED_SELECTION_MERGE_AUTO_MIN_SORTED_RATIO);
+ Preconditions.checkState(
+ _sortedSelectionMergeAutoMinSortedRatio >= 0 && _sortedSelectionMergeAutoMinSortedRatio <= 1,
+ "Invalid configuration: sortedSelectionMergeAutoMinSortedRatio: %s must be in [0, 1]",
+ _sortedSelectionMergeAutoMinSortedRatio);
+ _sortedSelectionMergeBlockSize = queryExecutorConfig.getProperty(Server.SORTED_SELECTION_MERGE_BLOCK_SIZE,
+ Server.DEFAULT_SORTED_SELECTION_MERGE_BLOCK_SIZE);
+ Preconditions.checkState(_sortedSelectionMergeBlockSize > 0,
+ "Invalid configuration: sortedSelectionMergeBlockSize: %s must be positive", _sortedSelectionMergeBlockSize);
LOGGER.info("Initialized plan maker with maxExecutionThreads: {}, defaultExecutionThreads: {}, "
+ "maxInitialResultHolderCapacity: {}, numGroupsLimit: {}, minSegmentGroupTrimSize: {}, "
- + "minServerGroupTrimSize: {}, groupByTrimThreshold: {}",
+ + "minServerGroupTrimSize: {}, groupByTrimThreshold: {}, sortedSelectionMergeAutoMinSortedRatio: {}, "
+ + "sortedSelectionMergeBlockSize: {}",
_maxExecutionThreads, _defaultExecutionThreads, _maxInitialResultHolderCapacity, _numGroupsLimit,
- _minSegmentGroupTrimSize, _minServerGroupTrimSize, _groupByTrimThreshold);
+ _minSegmentGroupTrimSize, _minServerGroupTrimSize, _groupByTrimThreshold,
+ _sortedSelectionMergeAutoMinSortedRatio, _sortedSelectionMergeBlockSize);
}
@VisibleForTesting
@@ -291,14 +307,14 @@ public class InstancePlanMakerImplV2 implements PlanMaker {
}
Double sortedSelectionMergeAutoMinSortedRatio =
QueryOptionsUtils.getSortedSelectionMergeAutoMinSortedRatio(queryOptions);
- if (sortedSelectionMergeAutoMinSortedRatio != null) {
- queryContext.setSortedSelectionMergeAutoMinSortedRatio(sortedSelectionMergeAutoMinSortedRatio);
- }
+ queryContext.setSortedSelectionMergeAutoMinSortedRatio(sortedSelectionMergeAutoMinSortedRatio != null
+ ? sortedSelectionMergeAutoMinSortedRatio
+ : _sortedSelectionMergeAutoMinSortedRatio);
Integer sortedSelectionMergeBlockSize =
QueryOptionsUtils.getSortedSelectionMergeBlockSize(queryOptions);
- if (sortedSelectionMergeBlockSize != null) {
- queryContext.setSortedSelectionMergeBlockSize(sortedSelectionMergeBlockSize);
- }
+ queryContext.setSortedSelectionMergeBlockSize(sortedSelectionMergeBlockSize != null
+ ? sortedSelectionMergeBlockSize
+ : _sortedSelectionMergeBlockSize);
}
// Set group-by query options
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java b/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java
index 384f800c8b2..7a5741826e6 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java
@@ -44,7 +44,6 @@ import org.apache.pinot.core.util.MemoizedClassAssociation;
import org.apache.pinot.segment.spi.datasource.DataSource;
import org.apache.pinot.spi.config.table.FieldConfig;
import org.apache.pinot.spi.data.Schema;
-import org.apache.pinot.spi.utils.CommonConstants.Broker;
import org.apache.pinot.spi.utils.CommonConstants.Server;
import org.apache.pinot.spi.utils.CommonConstants.Server.SortedSelectionMergeMode;
import org.slf4j.Logger;
@@ -148,7 +147,7 @@ public class QueryContext {
private double _sortedSelectionMergeAutoMinSortedRatio =
Server.DEFAULT_SORTED_SELECTION_MERGE_AUTO_MIN_SORTED_RATIO;
/// Output block size (rows) for the streaming selection ORDER BY combine
- private int _sortedSelectionMergeBlockSize = Broker.DEFAULT_SORTED_SELECTION_MERGE_BLOCK_SIZE;
+ private int _sortedSelectionMergeBlockSize = Server.DEFAULT_SORTED_SELECTION_MERGE_BLOCK_SIZE;
// Guards the one-time warning in isSortedSelectionMergeEnabled() so that an unresolved AUTO does not log once per
// segment/combine call for the same query.
private volatile boolean _unresolvedSortedSelectionMergeModeWarned;
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/CombineSlowOperatorsTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/CombineSlowOperatorsTest.java
index da476ee0dbd..786bebaf314 100644
--- a/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/CombineSlowOperatorsTest.java
+++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/CombineSlowOperatorsTest.java
@@ -179,9 +179,10 @@ public class CombineSlowOperatorsTest {
/// The merge loop drains its heap on the caller thread and, for single-block cursors, may never re-enter a child
/// operator, so it must check the query deadline itself. With an already-expired deadline the operator must surface a
- /// timeout <b>before</b> activating any child - asserted via {@code _operationInProgress}, which also keeps the test
- /// from passing vacuously if the timeout came from somewhere else.
- @Test
+ /// timeout **before** activating any child - asserted via `_operationInProgress`, which also keeps the test
+ /// from passing vacuously if the timeout came from somewhere else. A regression would drive a [SlowOperator], which
+ /// sleeps for an hour, so the test timeout turns that hang into a failure.
+ @Test(timeOut = 30_000)
public void testStreamingSelectionOrderByCombineOperatorHonorsDeadline() {
// A real latch, never awaited: this test asserts that no child is driven at all, so there is nothing to wait
// for. It is passed only so every getOperators() call site in this class looks alike; SlowOperator null-guards
@@ -221,7 +222,7 @@ public class CombineSlowOperatorsTest {
}
/// Submits the combine operator on a separate thread, waits for a child operator to start, cancels the future, and
- /// asserts the operator surfaces an {@link ExceptionResultsBlock}. When {@code operatorsToVerify} is non-empty
+ /// asserts the operator surfaces an [ExceptionResultsBlock]. When `operatorsToVerify` is non-empty
/// (the single-threaded streaming combine, whose merge runs the children on the caller thread), it additionally
/// asserts the cancellation was genuine - a child actually started and none completed normally - so the test
/// cannot pass vacuously on an unrelated failure or by never exercising the cancel path.
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java
index ae269cc1e76..606705708cc 100644
--- a/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java
+++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java
@@ -23,18 +23,25 @@ import java.io.IOException;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
+import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.stream.Collectors;
+import javax.annotation.Nullable;
import org.apache.commons.io.FileUtils;
+import org.apache.pinot.common.proto.Plan.ExplainNode.AttributeValue;
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.ExceptionResultsBlock;
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.plan.CombinePlanNode;
+import org.apache.pinot.core.plan.ExplainInfo;
import org.apache.pinot.core.plan.PlanNode;
import org.apache.pinot.core.plan.maker.InstancePlanMakerImplV2;
import org.apache.pinot.core.plan.maker.PlanMaker;
@@ -54,6 +61,8 @@ import org.apache.pinot.spi.config.table.TableType;
import org.apache.pinot.spi.data.FieldSpec;
import org.apache.pinot.spi.data.Schema;
import org.apache.pinot.spi.data.readers.GenericRow;
+import org.apache.pinot.spi.exception.QueryErrorCode;
+import org.apache.pinot.spi.exception.QueryErrorMessage;
import org.apache.pinot.spi.utils.CommonConstants.Server;
import org.apache.pinot.spi.utils.CommonConstants.Server.SortedSelectionMergeMode;
import org.apache.pinot.spi.utils.ReadMode;
@@ -71,22 +80,22 @@ import static org.testng.Assert.assertTrue;
import static org.testng.Assert.expectThrows;
-/// Combine-level tests for {@link StreamingSelectionOrderByCombineOperator} (step-3 operator) and its wiring into
-/// {@link CombinePlanNode#getCombineOperator()} (step-4).
+/// Combine-level tests for [StreamingSelectionOrderByCombineOperator] (step-3 operator) and its wiring into
+/// [CombinePlanNode#getCombineOperator()] (step-4).
///
-/// <p>The streaming combine must return the same globally-sorted top-K rows as the default
-/// {@link MinMaxValueBasedSelectionOrderByCombineOperator}, only (in streaming mode) spread across several bounded
-/// blocks. Each functional test therefore asserts <b>streaming-vs-non-streaming parity</b>: it runs the identical query
-/// twice over the same in-memory segments - once with {@code sortedSelectionMergeMode=ON} (asserting the new
-/// operator was actually selected) and once with the hint off (asserting the {@code MinMax} operator was selected) -
+/// The streaming combine must return the same globally-sorted top-K rows as the default
+/// [MinMaxValueBasedSelectionOrderByCombineOperator], only (in streaming mode) spread across several bounded
+/// blocks. Each functional test therefore asserts **streaming-vs-non-streaming parity**: it runs the identical query
+/// twice over the same in-memory segments - once with `sortedSelectionMergeMode=ON` (asserting the new
+/// operator was actually selected) and once with the hint off (asserting the `MinMax` operator was selected) -
/// then checks the two row sets are equal as a multiset and that the streaming output is fully sorted by the order-by
/// comparator.
///
-/// <p>To keep the top-K boundary unambiguous (operators may legitimately disagree on which of several rows that tie on
+/// To keep the top-K boundary unambiguous (operators may legitimately disagree on which of several rows that tie on
/// every order-by key fall inside the limit) every parity query ends its ORDER BY with the globally-unique
-/// {@code valCol}
+/// `valCol`
/// so the comparator is a total order; the merge still genuinely interleaves segments because the primary sort column
-/// ({@code sortedCol}) overlaps across segments. Multiset (rather than positional) comparison then tolerates only the
+/// (`sortedCol`) overlaps across segments. Multiset (rather than positional) comparison then tolerates only the
/// harmless reordering of fully-equal projected rows.
public class StreamingSelectionOrderByCombineOperatorTest {
private static final File TEMP_DIR =
@@ -101,6 +110,10 @@ public class StreamingSelectionOrderByCombineOperatorTest {
/// where MinMax emits INT), which an all-INT suite cannot observe.
private static final String LONG_COL = "longCol";
private static final String STR_COL = "strCol";
+ /// Takes the LIMIT. `tailCol` descends within each `_sortedSegments` segment, so it is not physically sorted and
+ /// each streaming child takes the run path: one block per distinct sortedCol value, which is one row per block.
+ private static final String ONE_ROW_PER_BLOCK_QUERY =
+ "SELECT sortedCol, tailCol, valCol FROM testTable ORDER BY sortedCol, tailCol LIMIT %d";
/// Create (MAX_NUM_THREADS_PER_QUERY * 2) sorted segments so the leaf runs plan nodes across multiple threads.
private static final int NUM_SEGMENTS = QueryMultiThreadingUtils.MAX_NUM_THREADS_PER_QUERY * 2;
@@ -452,7 +465,7 @@ public class StreamingSelectionOrderByCombineOperatorTest {
/// Correctness under single-column ties with an OFFSET: LIMIT 10 OFFSET 20 straddles the same sortedCol=0/1
/// boundary as [#testSingleColumnTieDeferralPreservesOrderByValuesAcrossLimitStraddle] (limit + offset = 30), but
/// exercises it with a nonzero offset. As [#testLimitOffsetParity] documents, the server (and this combine) retains
- /// {@code limit + offset} rows -- the offset is trimmed by the broker afterwards -- so 30, not 10, rows come back
+ /// `limit + offset` rows -- the offset is trimmed by the broker afterwards -- so 30, not 10, rows come back
/// here too; what differs from the straddle test is only the query shape, not the row count.
@Test
public void testSingleColumnTieDeferralPreservesOrderByValuesAcrossOffset() {
@@ -617,6 +630,9 @@ public class StreamingSelectionOrderByCombineOperatorTest {
+ result._numDocsScanned + " of " + totalDocs);
assertTrue(result._numDocsScanned <= NUM_RECORDS_PER_SEGMENT,
"Only the lowest-range segment should be scanned, but docs scanned was: " + result._numDocsScanned);
+ // The pruning is visible in the stats only as the gap between these two: see the class javadoc.
+ assertEquals(result._numSegmentsProcessed, NUM_SEGMENTS, "numSegmentsProcessed counts every segment");
+ assertEquals(result._numSegmentsMatched, 1, "numSegmentsMatched must exclude never-activated segments");
}
/// A cursor activated *late* can hold a smaller row than the retained leader. Two things must be right for it to
@@ -756,6 +772,219 @@ public class StreamingSelectionOrderByCombineOperatorTest {
expectThrows(IllegalStateException.class, streamingCombine::processSegments);
}
+ // Segment acquire/release lifecycle. The merge acquires a segment when it activates the segment's cursor, so every
+ // way out of the merge has to release whatever is still held; a missed release pins the segment for the lifetime of
+ // the server. `_sortedSegments` segment i covers sortedCol [i * 50, i * 50 + 100).
+
+ @Test
+ public void testFullDrainReleasesEverySegment() {
+ QueryContext queryContext =
+ hintedContext("SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 100000");
+ List<InstrumentedSegmentOperator> children = instrument(_sortedSegments, queryContext);
+ List<BaseResultsBlock> blocks = drain(combineOver(children, queryContext));
+ assertEquals(rowsOf(blocks).size(), NUM_SEGMENTS * NUM_RECORDS_PER_SEGMENT);
+ assertEquals(numAcquired(children), NUM_SEGMENTS, "A full drain must activate every segment");
+ assertEveryAcquireReleased(children);
+ }
+
+ @Test
+ public void testLimitReachedReleasesSegmentsStillHeld() {
+ // 60 rows end around sortedCol 55, where segments 0 and 1 are both active and neither is drained, so finish()
+ // is the only thing that can release them. Segment 2 starts at 100 and must never be acquired.
+ QueryContext queryContext =
+ hintedContext("SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 60");
+ queryContext.setSortedSelectionMergeBlockSize(1000);
+ List<InstrumentedSegmentOperator> children = instrument(_sortedSegments, queryContext);
+ List<BaseResultsBlock> blocks = drain(combineOver(children, queryContext));
+ assertEquals(rowsOf(blocks).size(), 60);
+ assertEquals(numAcquired(children), 2, "Only the two segments overlapping the first 60 rows may be acquired");
+ assertEveryAcquireReleased(children);
+ }
+
+ @Test
+ public void testStopReleasesSegmentsHeldMidMerge() {
+ QueryContext queryContext =
+ hintedContext("SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 200");
+ queryContext.setSortedSelectionMergeBlockSize(60);
+ List<InstrumentedSegmentOperator> children = instrument(_sortedSegments, queryContext);
+ StreamingSelectionOrderByCombineOperator combine = combineOver(children, queryContext);
+
+ BaseResultsBlock first = combine.nextBlock();
+ assertTrue(first instanceof SelectionResultsBlock, "Expected a data block, got: " + first.getClass());
+ assertTrue(numHeld(children) > 0, "Segments must be held mid-merge, else the test is vacuous");
+
+ combine.stop();
+ assertEveryAcquireReleased(children);
+ // A driver may stop more than once; release must not run again for a segment already released.
+ combine.stop();
+ assertEveryAcquireReleased(children);
+ assertTrue(combine.nextBlock() instanceof MetadataResultsBlock, "A stopped merge must not resume");
+ assertEveryAcquireReleased(children);
+ }
+
+ @Test
+ public void testChildExceptionMidMergeReleasesEverySegment() {
+ // Segment 0 fails on its 60th block (sortedCol 59), after segment 1 (which starts at 50) has been activated, so a
+ // segment other than the failing one is also held when the exception unwinds the merge.
+ QueryContext queryContext = hintedContext(String.format(ONE_ROW_PER_BLOCK_QUERY, 200));
+ queryContext.setSortedSelectionMergeBlockSize(1000);
+ List<InstrumentedSegmentOperator> children = instrument(_sortedSegments, queryContext);
+ children.get(0)._interceptor = (blockNumber, block) -> {
+ if (blockNumber == 60) {
+ throw new IllegalStateException("Injected child failure");
+ }
+ return block;
+ };
+ BaseResultsBlock result = combineOver(children, queryContext).nextBlock();
+ assertTrue(result instanceof ExceptionResultsBlock, "Expected an exception block, got: " + result.getClass());
+ assertEquals(numAcquired(children), 2, "Segments 0 and 1 must both be held when the failure hits");
+ assertEveryAcquireReleased(children);
+ }
+
+ @Test
+ public void testChildErrorMidMergeReleasesEverySegment() {
+ // An Error is rethrown rather than turned into an exception block, so single-stage execution has no stop() after
+ // it: the merge's own Throwable handler is the only release.
+ QueryContext queryContext = hintedContext(String.format(ONE_ROW_PER_BLOCK_QUERY, 200));
+ queryContext.setSortedSelectionMergeBlockSize(1000);
+ List<InstrumentedSegmentOperator> children = instrument(_sortedSegments, queryContext);
+ children.get(0)._interceptor = (blockNumber, block) -> {
+ if (blockNumber == 60) {
+ throw new OutOfMemoryError("Injected child error");
+ }
+ return block;
+ };
+ StreamingSelectionOrderByCombineOperator combine = combineOver(children, queryContext);
+ expectThrows(OutOfMemoryError.class, combine::nextBlock);
+ assertEquals(numAcquired(children), 2, "Segments 0 and 1 must both be held when the error hits");
+ assertEveryAcquireReleased(children);
+ }
+
+ /// A block whose schema disagrees with the merge's is dropped together with the rest of its segment, and reported
+ /// once as a MERGE_RESPONSE error however many segments disagree the same way. A segment's schema does not change
+ /// between its own blocks in practice (a reload swaps the whole segment); diverging on a middle block here is what
+ /// separates "drop the rest of the segment" from "drop one block".
+ @Test
+ public void testSchemaMismatchDropsTheRestOfTheSegmentAndIsReportedOnce() {
+ String query = String.format(ONE_ROW_PER_BLOCK_QUERY, 100000);
+ QueryContext queryContext = hintedContext(query);
+ queryContext.setSortedSelectionMergeBlockSize(7);
+ List<InstrumentedSegmentOperator> children = instrument(_sortedSegments, queryContext);
+ int divergingBlock = 11;
+ for (int segment : List.of(1, 2)) {
+ children.get(segment)._interceptor = (blockNumber, block) -> blockNumber == divergingBlock
+ ? withMismatchedSchema((SelectionResultsBlock) block)
+ : block;
+ }
+
+ List<BaseResultsBlock> blocks = drain(combineOver(children, queryContext));
+ assertTrue(blocks.get(blocks.size() - 1) instanceof MetadataResultsBlock,
+ "A schema mismatch must not fail the query, got: " + blocks.get(blocks.size() - 1).getClass());
+ List<Object[]> rows = rowsOf(blocks);
+ assertSorted(rows, orderByComparator(query, false));
+
+ // One row per block, so each diverging segment contributes exactly the rows of the blocks before the mismatch.
+ int valColIndex = List.of(blocks.get(0).getDataSchema().getColumnNames()).indexOf(VAL_COL);
+ int[] rowsPerSegment = new int[NUM_SEGMENTS];
+ for (Object[] row : rows) {
+ int valCol = (Integer) row[valColIndex];
+ int segment = valCol / 1_000_000;
+ rowsPerSegment[segment]++;
+ if (segment == 1 || segment == 2) {
+ assertTrue(valCol % 1_000_000 < divergingBlock - 1,
+ "Segment " + segment + " emitted a row from at or after the mismatched block: " + valCol);
+ }
+ }
+ for (int segment = 0; segment < NUM_SEGMENTS; segment++) {
+ int expected = segment == 1 || segment == 2 ? divergingBlock - 1 : NUM_RECORDS_PER_SEGMENT;
+ assertEquals(rowsPerSegment[segment], expected, "Unexpected row count for segment " + segment);
+ }
+ for (int segment : List.of(1, 2)) {
+ assertEquals(children.get(segment)._numBlocks, divergingBlock,
+ "Segment " + segment + " must not be read past the mismatched block");
+ }
+ assertEveryAcquireReleased(children);
+
+ // The same mismatch in two segments produces one message, attached to one block rather than re-attached to every
+ // block streamed after it.
+ List<QueryErrorMessage> errors = new ArrayList<>();
+ for (BaseResultsBlock block : blocks) {
+ if (block.getErrorMessages() != null) {
+ errors.addAll(block.getErrorMessages());
+ }
+ }
+ assertEquals(errors.size(), 1, "Expected exactly one mismatch error, got: " + errors);
+ assertEquals(errors.get(0).getErrCode(), QueryErrorCode.MERGE_RESPONSE);
+ }
+
+ /// Returns `block` with its last column retyped from INT to LONG, the shape a newly reloaded segment takes when a
+ /// column's type changes.
+ private static SelectionResultsBlock withMismatchedSchema(SelectionResultsBlock block) {
+ DataSchema dataSchema = block.getDataSchema();
+ ColumnDataType[] columnDataTypes = dataSchema.getColumnDataTypes().clone();
+ int last = columnDataTypes.length - 1;
+ assertEquals(columnDataTypes[last], ColumnDataType.INT);
+ columnDataTypes[last] = ColumnDataType.LONG;
+ return new SelectionResultsBlock(new DataSchema(dataSchema.getColumnNames(), columnDataTypes), block.getRows(),
+ block.getComparator(), block.getQueryContext());
+ }
+
+ /// Plans one [InstrumentedSegmentOperator] per segment, in segment order, the way the prefetch path wraps the real
+ /// streaming leaf.
+ private static List<InstrumentedSegmentOperator> instrument(List<IndexSegment> segments,
+ QueryContext queryContext) {
+ List<InstrumentedSegmentOperator> children = new ArrayList<>(segments.size());
+ for (IndexSegment segment : segments) {
+ children.add(new InstrumentedSegmentOperator(
+ PLAN_MAKER.makeStreamingSegmentPlanNode(new SegmentContext(segment), queryContext), segment));
+ }
+ return children;
+ }
+
+ private static StreamingSelectionOrderByCombineOperator combineOver(List<InstrumentedSegmentOperator> children,
+ QueryContext queryContext) {
+ return new StreamingSelectionOrderByCombineOperator(new ArrayList<>(children), queryContext, EXECUTOR);
+ }
+
+ /// Drives `combine` until it returns something other than a data block, and returns every block, that one last.
+ private static List<BaseResultsBlock> drain(Operator<?> combine) {
+ List<BaseResultsBlock> blocks = new ArrayList<>();
+ while (true) {
+ BaseResultsBlock block = (BaseResultsBlock) combine.nextBlock();
+ blocks.add(block);
+ if (!(block instanceof SelectionResultsBlock)) {
+ return blocks;
+ }
+ assertTrue(blocks.size() < 1_000_000, "Streaming combine did not terminate");
+ }
+ }
+
+ private static List<Object[]> rowsOf(List<BaseResultsBlock> blocks) {
+ List<Object[]> rows = new ArrayList<>();
+ for (BaseResultsBlock block : blocks) {
+ if (block instanceof SelectionResultsBlock) {
+ rows.addAll(((SelectionResultsBlock) block).getRows());
+ }
+ }
+ return rows;
+ }
+
+ private static int numAcquired(List<InstrumentedSegmentOperator> children) {
+ return children.stream().mapToInt(child -> child._numAcquires).sum();
+ }
+
+ private static int numHeld(List<InstrumentedSegmentOperator> children) {
+ return children.stream().mapToInt(child -> child._numAcquires - child._numReleases).sum();
+ }
+
+ private static void assertEveryAcquireReleased(List<InstrumentedSegmentOperator> children) {
+ for (int i = 0; i < children.size(); i++) {
+ InstrumentedSegmentOperator child = children.get(i);
+ assertTrue(child._numAcquires <= 1, "Segment " + i + " acquired " + child._numAcquires + " times");
+ assertEquals(child._numReleases, child._numAcquires, "Segment " + i + " acquires and releases differ");
+ }
+ }
+
@Test
public void testHintOffSelectsMinMaxOperator() {
// Default behavior is unchanged when the hint is off: the classic MinMax operator is still selected.
@@ -815,7 +1044,65 @@ public class StreamingSelectionOrderByCombineOperatorTest {
+ combineOperator.getClass().getSimpleName());
}
- /// Asserts streaming-vs-blocking parity for {@code query}. Runs the MinMax combine on the blocking path (hint off)
+ @Test
+ public void testExplainAttributes() {
+ @Language("sql") String query = "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50";
+ QueryContext queryContext = hintedContext(query);
+ queryContext.setSortedSelectionMergeBlockSize(7);
+ Operator<?> combineOperator = planCombineOperator(_sortedSegments, queryContext, true);
+ ExplainInfo explainInfo = combineOperator.getExplainInfo();
+ assertEquals(explainInfo.getTitle(), "CombineSelectOrderbyStreaming");
+ Map<String, AttributeValue> attributes = explainInfo.getAttributes();
+ assertEquals(attributes.get("blockSize").getLong(), 7);
+ // Query-wide settings must not be summed when the broker merges this node across workers.
+ assertEquals(attributes.get("blockSize").getMergeType(), AttributeValue.MergeType.IDEMPOTENT);
+ assertTrue(attributes.get("frontierPruning").getBool());
+ assertFalse(attributes.get("deferTiedCursors").getBool(), "Two order-by expressions must not defer ties");
+ assertEquals(attributes.get("numSegments").getLong(), NUM_SEGMENTS);
+ assertEquals(attributes.get("numSortedSegments").getLong(), NUM_SEGMENTS);
+
+ // Explaining must not start the merge: the same operator still returns the full, correct result afterwards.
+ List<Object[]> rows = new ArrayList<>();
+ while (true) {
+ BaseResultsBlock block = (BaseResultsBlock) combineOperator.nextBlock();
+ if (block instanceof MetadataResultsBlock) {
+ break;
+ }
+ rows.addAll(((SelectionResultsBlock) block).getRows());
+ }
+ assertMultisetEquals(rows, run(_sortedSegments, query, false, false, false, 0)._rows);
+ }
+
+ @Test
+ public void testExplainAttributesReflectTheQueryShape() {
+ assertTrue(explainAttributes(_sortedSegments, "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol LIMIT 5",
+ false).get("deferTiedCursors").getBool(), "A single order-by expression defers ties");
+ assertFalse(explainAttributes(_sortedSegments,
+ "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 5", true)
+ .get("frontierPruning").getBool(), "Null handling disables frontier pruning");
+ assertEquals(explainAttributes(_mixedSegments,
+ "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 5", false)
+ .get("numSortedSegments").getLong(), 2);
+ assertEquals(explainAttributes(_unsortedSegments,
+ "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 5", false)
+ .get("numSortedSegments").getLong(), 0);
+
+ Map<String, AttributeValue> expression = explainAttributes(_sortedSegments,
+ "SELECT sortedCol, valCol FROM testTable ORDER BY ADD(sortedCol, 1), valCol LIMIT 5", false);
+ assertFalse(expression.get("frontierPruning").getBool(), "An expression leading order-by has no min/max to prune");
+ assertFalse(expression.containsKey("numSortedSegments"), "Sortedness is undefined for an expression");
+ }
+
+ private static Map<String, AttributeValue> explainAttributes(List<IndexSegment> segments,
+ @Language("sql") String query, boolean nullHandling) {
+ QueryContext queryContext = hintedContext(query);
+ queryContext.setNullHandlingEnabled(nullHandling);
+ Operator<?> combineOperator = planCombineOperator(segments, queryContext, true);
+ assertTrue(combineOperator instanceof StreamingSelectionOrderByCombineOperator);
+ return combineOperator.getExplainInfo().getAttributes();
+ }
+
+ /// Asserts streaming-vs-blocking parity for `query`. Runs the MinMax combine on the blocking path (hint off)
/// as the reference, then runs the streaming combine on the streaming path, asserting it selects the streaming
/// operator and produces rows that are sorted by the order-by comparator and equal the MinMax rows as a multiset,
/// plus the bounded-flush invariants. A small block size forces several bounded data blocks before the metadata
@@ -852,7 +1139,7 @@ public class StreamingSelectionOrderByCombineOperatorTest {
}
}
- /// Plans the combine operator that would run {@code queryContext} over {@code segments}, without driving it. Use
+ /// Plans the combine operator that would run `queryContext` over `segments`, without driving it. Use
/// this for gate assertions: an operator that is not the streaming combine (the fallbacks) does not terminate with
/// a metadata block, so it cannot be driven by the streaming loop in [#run].
private static Operator<?> planCombineOperator(List<IndexSegment> segments, QueryContext queryContext,
@@ -898,17 +1185,17 @@ public class StreamingSelectionOrderByCombineOperatorTest {
@Test
public void testAutoHonoursTheMinSortedRatioThreshold() {
// _mixedSegments is 2 sorted of 4, so a ratio of exactly 0.5 must pass and anything above it must fail. This pins
- // the comparison as >= rather than >, and pins that the threshold is read from the query context.
+ // the comparison as >= rather than >, and pins that the threshold is read from the query option.
String query = "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50";
- QueryContext atThreshold = modeContext(query, SortedSelectionMergeMode.AUTO);
- atThreshold.setSortedSelectionMergeAutoMinSortedRatio(0.5);
+ QueryContext atThreshold =
+ modeContext("SET sortedSelectionMergeAutoMinSortedRatio=0.5; " + query, SortedSelectionMergeMode.AUTO);
makeStreamingInstancePlan(_mixedSegments, atThreshold);
assertEquals(atThreshold.getSortedSelectionMergeMode(), SortedSelectionMergeMode.ON,
"A sorted ratio equal to the threshold must select the streaming merge");
- QueryContext aboveThreshold = modeContext(query, SortedSelectionMergeMode.AUTO);
- aboveThreshold.setSortedSelectionMergeAutoMinSortedRatio(0.75);
+ QueryContext aboveThreshold =
+ modeContext("SET sortedSelectionMergeAutoMinSortedRatio=0.75; " + query, SortedSelectionMergeMode.AUTO);
makeStreamingInstancePlan(_mixedSegments, aboveThreshold);
assertEquals(aboveThreshold.getSortedSelectionMergeMode(), SortedSelectionMergeMode.OFF,
"A sorted ratio below the threshold must keep the MinMax combine");
@@ -1018,14 +1305,14 @@ public class StreamingSelectionOrderByCombineOperatorTest {
String query = "SET allowReverseOrder=true; SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol DESC, "
+ "valCol DESC LIMIT 50";
- QueryContext atThreshold = modeContext(query, SortedSelectionMergeMode.AUTO);
- atThreshold.setSortedSelectionMergeAutoMinSortedRatio(0.5);
+ QueryContext atThreshold =
+ modeContext("SET sortedSelectionMergeAutoMinSortedRatio=0.5; " + query, SortedSelectionMergeMode.AUTO);
makeStreamingInstancePlan(_mixedSegments, atThreshold);
assertEquals(atThreshold.getSortedSelectionMergeMode(), SortedSelectionMergeMode.ON,
"A DESC sorted ratio equal to the threshold must select the streaming merge");
- QueryContext aboveThreshold = modeContext(query, SortedSelectionMergeMode.AUTO);
- aboveThreshold.setSortedSelectionMergeAutoMinSortedRatio(0.75);
+ QueryContext aboveThreshold =
+ modeContext("SET sortedSelectionMergeAutoMinSortedRatio=0.75; " + query, SortedSelectionMergeMode.AUTO);
makeStreamingInstancePlan(_mixedSegments, aboveThreshold);
assertEquals(aboveThreshold.getSortedSelectionMergeMode(), SortedSelectionMergeMode.OFF,
"A DESC sorted ratio below the threshold must keep the MinMax combine");
@@ -1069,7 +1356,7 @@ public class StreamingSelectionOrderByCombineOperatorTest {
false);
}
- /// Runs the streaming instance plan for its side effect on {@code queryContext}: resolving AUTO from segment
+ /// Runs the streaming instance plan for its side effect on `queryContext`: resolving AUTO from segment
/// metadata. The plan itself is discarded; the gate tests assert on the resolved mode and on the operator that
/// [#planCombineOperator] then selects.
private static void makeStreamingInstancePlan(List<IndexSegment> segments, QueryContext queryContext) {
@@ -1081,7 +1368,7 @@ public class StreamingSelectionOrderByCombineOperatorTest {
});
}
- /// Builds a query context for {@code query} with the streaming merge mode set to {@code mode}.
+ /// Builds a query context for `query` with the streaming merge mode set to `mode`.
private static QueryContext modeContext(@Language("sql") String query, SortedSelectionMergeMode mode) {
QueryContext queryContext = QueryContextConverterUtils.getQueryContext(query);
queryContext.setSortedSelectionMergeMode(mode);
@@ -1089,7 +1376,7 @@ public class StreamingSelectionOrderByCombineOperatorTest {
return queryContext;
}
- /// Builds a query context for {@code query} with the streaming merge hint on.
+ /// Builds a query context for `query` with the streaming merge hint on.
private static QueryContext hintedContext(@Language("sql") String query) {
QueryContext queryContext = QueryContextConverterUtils.getQueryContext(query);
queryContext.setSortedSelectionMergeMode(SortedSelectionMergeMode.ON);
@@ -1097,7 +1384,7 @@ public class StreamingSelectionOrderByCombineOperatorTest {
return queryContext;
}
- /// Runs one combine over {@code segments} and collects its rows, blocks, schema and docs-scanned stat.
+ /// Runs one combine over `segments` and collects its rows, blocks, schema and docs-scanned stat.
private Result run(List<IndexSegment> segments, @Language("sql") String query, boolean hintOn, boolean nullHandling,
boolean streaming, int blockSize) {
QueryContext queryContext = QueryContextConverterUtils.getQueryContext(query);
@@ -1126,6 +1413,7 @@ public class StreamingSelectionOrderByCombineOperatorTest {
}
result._numDocsScanned = block.getNumDocsScanned();
result._numSegmentsMatched = block.getNumSegmentsMatched();
+ result._numSegmentsProcessed = block.getNumSegmentsProcessed();
break;
}
SelectionResultsBlock dataBlock = (SelectionResultsBlock) block;
@@ -1173,7 +1461,7 @@ public class StreamingSelectionOrderByCombineOperatorTest {
/// Canonicalizes rows for multiset comparison. Each cell is encoded with its runtime class so a stored-type / boxing
/// regression (e.g. a LONG emitted where the reference emits INT) changes the encoding and fails the assertion, which
- /// a plain {@code Arrays.toString} (type-blind) comparison would miss.
+ /// a plain `Arrays.toString` (type-blind) comparison would miss.
private static List<String> toCanonical(List<Object[]> rows) {
return rows.stream().map(row -> {
StringBuilder sb = new StringBuilder("[");
@@ -1218,5 +1506,44 @@ public class StreamingSelectionOrderByCombineOperatorTest {
private int _numBlocks;
private long _numDocsScanned;
private int _numSegmentsMatched;
+ private int _numSegmentsProcessed;
+ }
+
+ /// A segment child wrapped the way the prefetch path wraps it, counting acquire/release calls and optionally
+ /// replacing or failing a numbered child block. Acquire and release are counted rather than forwarded, since these
+ /// immutable test segments have nothing to fetch.
+ private static class InstrumentedSegmentOperator extends AcquireReleaseColumnsSegmentOperator {
+ private int _numAcquires;
+ private int _numReleases;
+ private int _numBlocks;
+ @Nullable
+ private BlockInterceptor _interceptor;
+
+ InstrumentedSegmentOperator(PlanNode planNode, IndexSegment indexSegment) {
+ super(planNode, indexSegment, null);
+ }
+
+ @Override
+ public void acquire() {
+ _numAcquires++;
+ }
+
+ @Override
+ public void release() {
+ _numReleases++;
+ }
+
+ @Override
+ protected BaseResultsBlock getNextBlock() {
+ BaseResultsBlock block = super.getNextBlock();
+ _numBlocks++;
+ return _interceptor != null ? _interceptor.intercept(_numBlocks, block) : block;
+ }
+ }
+
+ @FunctionalInterface
+ private interface BlockInterceptor {
+ /// Returns the block to hand the merge in place of `block`, the child's `blockNumber`-th (1-based).
+ BaseResultsBlock intercept(int blockNumber, BaseResultsBlock block);
}
}
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperatorTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperatorTest.java
index f77bdc4d4b4..62a83bc7816 100644
--- a/pinot-core/src/test/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperatorTest.java
+++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperatorTest.java
@@ -73,20 +73,20 @@ import static org.testng.Assert.assertNull;
import static org.testng.Assert.assertTrue;
-/// Segment-level tests for {@link StreamingSelectionOrderByOperator}.
+/// Segment-level tests for [StreamingSelectionOrderByOperator].
///
-/// <p>The operator emits the same globally-sorted rows as the existing materialized selection ORDER BY operators, only
-/// spread across many lazily-produced blocks. Each test therefore asserts <b>stream-vs-materialized parity</b>: it
-/// drives the streaming operator through {@link SelectionPlanNode} with {@code sortedSelectionMergeMode=ON},
+/// The operator emits the same globally-sorted rows as the existing materialized selection ORDER BY operators, only
+/// spread across many lazily-produced blocks. Each test therefore asserts **stream-vs-materialized parity**: it
+/// drives the streaming operator through [SelectionPlanNode] with `sortedSelectionMergeMode=ON`,
/// concatenates
-/// every {@link Operator#nextBlock()} output until {@code null}, and asserts the concatenation equals the single block
-/// the materialized operator ({@link SelectionPartiallyOrderedByLinearOperator} / {@link SelectionOrderByOperator})
+/// every [Operator#nextBlock()] output until `null`, and asserts the concatenation equals the single block
+/// the materialized operator ([SelectionPartiallyOrderedByLinearOperator] / [SelectionOrderByOperator])
/// produces for the identical query with the hint off.
///
-/// <p>To keep element-wise comparison deterministic (priority-queue draining is not stable for rows that tie on every
-/// order-by column) each fixture makes the order-by column tuple unique per row: the {@code _segment} fixture has a
-/// unique sorted column, while the {@code _dupSegment} / {@code _largeSegment} fixtures repeat the sorted column but
-/// pair it with a unique {@code TAIL_COL} order-by tail. The all-order-by-columns-tie case (where stream and
+/// To keep element-wise comparison deterministic (priority-queue draining is not stable for rows that tie on every
+/// order-by column) each fixture makes the order-by column tuple unique per row: the `_segment` fixture has a
+/// unique sorted column, while the `_dupSegment` / `_largeSegment` fixtures repeat the sorted column but
+/// pair it with a unique `TAIL_COL` order-by tail. The all-order-by-columns-tie case (where stream and
/// materialized output may legitimately differ in row order) is intentionally out of scope here and is covered by the
/// combine-level test via multiset comparison.
public class StreamingSelectionOrderByOperatorTest {
@@ -202,7 +202,7 @@ public class StreamingSelectionOrderByOperatorTest {
return records;
}
- /// Appends one run of {@code runSize} rows all sharing {@code sortedValue}, with a tail that descends within the run
+ /// Appends one run of `runSize` rows all sharing `sortedValue`, with a tail that descends within the run
/// (so the tail column is not globally sorted) and is unique per (sortedCol, tailCol) tuple.
private static void appendRun(List<GenericRow> records, int sortedValue, int runSize) {
for (int j = 0; j < runSize; j++) {
@@ -349,7 +349,7 @@ public class StreamingSelectionOrderByOperatorTest {
"Expected no nulls in the output when null handling is disabled");
}
- /// No-tail mode, single phase: {@link StreamingSelectionOrderByOperator} must skip a zero-document project block
+ /// No-tail mode, single phase: [StreamingSelectionOrderByOperator] must skip a zero-document project block
/// rather than read it as end-of-segment.
@Test
public void testEmptyProjectBlockDoesNotTruncateSortedScan() {
@@ -357,14 +357,14 @@ public class StreamingSelectionOrderByOperatorTest {
1, 0);
}
- /// Same, two phase: this is also the only path that asks the injected block for {@code getDocIds()}.
+ /// Same, two phase: this is also the only path that asks the injected block for `getDocIds()`.
@Test
public void testEmptyProjectBlockDoesNotTruncateTwoPhaseSortedScan() {
assertEmptyProjectBlocksAreSkipped(_segment, "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol LIMIT 25",
1, 2, 1, 0);
}
- /// Tail mode: the run scan reaches the project operator through {@code nextRow()}, which already tolerated empty
+ /// Tail mode: the run scan reaches the project operator through `nextRow()`, which already tolerated empty
/// blocks. Pins that the two scan paths agree.
@Test
public void testEmptyProjectBlockDoesNotTruncateRunScan() {
@@ -380,7 +380,7 @@ public class StreamingSelectionOrderByOperatorTest {
}
/// Empty blocks arriving immediately before exhaustion exercise the loop's other exit: the skip must fall through
- /// to the project operator's {@code null} rather than spin or emit a phantom row. The limit exceeds the fixture so
+ /// to the project operator's `null` rather than spin or emit a phantom row. The limit exceeds the fixture so
/// the scan actually reaches the end instead of stopping on the row budget.
@Test
public void testEmptyProjectBlocksBeforeExhaustionEndTheScan() {
@@ -388,14 +388,14 @@ public class StreamingSelectionOrderByOperatorTest {
0, 2);
}
- /// The skip loop sits directly on top of the {@code limit + offset} budget bookkeeping, so cover a non-zero offset.
+ /// The skip loop sits directly on top of the `limit + offset` budget bookkeeping, so cover a non-zero offset.
@Test
public void testEmptyProjectBlockIsSkippedWithOffset() {
assertEmptyProjectBlocksAreSkipped(_segment,
"SELECT sortedCol FROM testTable ORDER BY sortedCol LIMIT 15 OFFSET 10", 1, 2, 1, 0);
}
- /// {@code nextRow()} rebuilds the phase-1 null bitmaps in the same branch that pulls the next non-empty block, so
+ /// `nextRow()` rebuilds the phase-1 null bitmaps in the same branch that pulls the next non-empty block, so
/// the skip must not desynchronise a bitmap from the block it was built against. Uses the fixture that carries a
/// null in the order-by tail column, with null handling on.
@Test
@@ -489,7 +489,7 @@ public class StreamingSelectionOrderByOperatorTest {
}
}
- /// Each {@link ProjectPlanNode} build looks up one data source per projected column, so a plan that builds the
+ /// Each [ProjectPlanNode] build looks up one data source per projected column, so a plan that builds the
/// project twice looks up more than the materialized plan, which builds it once.
@Test
public void testDescIncompatibleFallbackBuildsTheProjectOnce() {
@@ -504,7 +504,7 @@ public class StreamingSelectionOrderByOperatorTest {
"The DESC-incompatible fall-through built the project twice");
}
- /// Plans {@code query} in {@code mode} against a segment counting data source lookups into {@code numLookups},
+ /// Plans `query` in `mode` against a segment counting data source lookups into `numLookups`,
/// asserting it lands on the materialized DESC operator.
private Operator<SelectionResultsBlock> planWithCountedDataSourceLookups(@Language("sql") String query,
SortedSelectionMergeMode mode, AtomicInteger numLookups) {
@@ -525,9 +525,9 @@ public class StreamingSelectionOrderByOperatorTest {
}
/// Every block the tail-to-sort path emits carries a structurally mutable row list. Defensive rather than a fix for
- /// a reachable failure: no consumer of this operator adds to the list today, but {@link SelectionResultsBlock} is a
- /// shared type and {@code SelectionOperatorUtils.mergeWithoutOrdering()} adds to the row list of the block it merges
- /// into, which a fixed-size {@code Arrays.asList} view would reject.
+ /// a reachable failure: no consumer of this operator adds to the list today, but [SelectionResultsBlock] is a
+ /// shared type and `SelectionOperatorUtils.mergeWithoutOrdering()` adds to the row list of the block it merges
+ /// into, which a fixed-size `Arrays.asList` view would reject.
@Test
public void testRunPathEmitsGrowableRowLists() {
QueryContext queryContext = QueryContextConverterUtils.getQueryContext(
@@ -570,7 +570,7 @@ public class StreamingSelectionOrderByOperatorTest {
}
/// A segment matching no rows must still emit exactly one block before signalling exhaustion: empty, but carrying
- /// the same {@link DataSchema} the identical query produces when it does match rows. Consumers therefore never have
+ /// the same [DataSchema] the identical query produces when it does match rows. Consumers therefore never have
/// to reconstruct a schema the segment already knows.
private void assertZeroMatchEmitsOneSchemaBlock(IndexSegment segment, @Language("sql") String zeroMatchQuery,
@Language("sql") String matchingQuery, int numSortedExpressions, boolean nullHandling) {
@@ -592,11 +592,11 @@ public class StreamingSelectionOrderByOperatorTest {
assertNull(operator.nextBlock(), "Exactly one block may precede exhaustion");
}
- /// Drives the operator twice over {@code segment}: once against the real project operator, once against one that
+ /// Drives the operator twice over `segment`: once against the real project operator, once against one that
/// splices zero-document blocks into its output, and asserts the emitted rows are identical.
///
- /// <p>No doc-id-set operator produces an empty block today, so the case has to be injected. The project operator is
- /// built with a deliberately small {@code maxDocsPerCall} so the scan spans several blocks.
+ /// No doc-id-set operator produces an empty block today, so the case has to be injected. The project operator is
+ /// built with a deliberately small `maxDocsPerCall` so the scan spans several blocks.
///
/// @param injectBeforeRealBlock 1-based index of the real block to splice empties in front of; 0 for none
/// @param numEmptyBlocks how many consecutive empty blocks to splice in at that point
@@ -625,7 +625,7 @@ public class StreamingSelectionOrderByOperatorTest {
}
}
- /// Builds a {@link StreamingSelectionOrderByOperator} directly (rather than through {@link SelectionPlanNode}) so
+ /// Builds a [StreamingSelectionOrderByOperator] directly (rather than through [SelectionPlanNode]) so
/// the project operator can be wrapped. Passing 0 for every injection parameter yields the undecorated operator.
private StreamingSelectionOrderByOperator buildOperatorWithInjectableProject(IndexSegment segment,
@Language("sql") String query, int numSortedExpressions, boolean nullHandling, int injectBeforeRealBlock,
@@ -665,10 +665,10 @@ public class StreamingSelectionOrderByOperatorTest {
}
/// Forwards everything to a real project operator, but splices zero-document blocks into its output:
- /// {@code numEmptyBlocks} of them just before the {@code injectBeforeRealBlock}-th real block, and
- /// {@code numTrailingEmptyBlocks} after the last real block but before exhaustion.
+ /// `numEmptyBlocks` of them just before the `injectBeforeRealBlock`-th real block, and
+ /// `numTrailingEmptyBlocks` after the last real block but before exhaustion.
///
- /// <p>Every real block is still delivered, merely deferred, so the decorator provably drops nothing of its own -
+ /// Every real block is still delivered, merely deferred, so the decorator provably drops nothing of its own -
/// any row loss observed by a test is the operator under test truncating its scan.
private static class EmptyBlockInjectingProjectOperator extends BaseProjectOperator<ValueBlock> {
private final BaseProjectOperator<?> _delegate;
@@ -791,13 +791,13 @@ public class StreamingSelectionOrderByOperatorTest {
}
}
- /// Runs {@code query} twice over {@code segment} - once with the streaming hint on, once off - and asserts the
+ /// Runs `query` twice over `segment` - once with the streaming hint on, once off - and asserts the
/// concatenated streaming blocks equal the materialized operator's single block, cell by cell. Returns the (now
/// exhausted) streaming operator so callers can make extra assertions on its execution statistics.
///
/// @param expectedMinBlocks the minimum number of non-null blocks the streaming operator must emit. Tail-mode
- /// cases that span multiple runs pass {@code >= 2} to prove the output is genuinely streamed; cases whose
- /// result fits in a single trimmed block pass {@code 1}.
+ /// cases that span multiple runs pass `>= 2` to prove the output is genuinely streamed; cases whose
+ /// result fits in a single trimmed block pass `1`.
private StreamingSelectionOrderByOperator assertParity(IndexSegment segment, @Language("sql") String query,
boolean nullHandling, int expectedMinBlocks) {
// Streaming path.
@@ -842,7 +842,7 @@ public class StreamingSelectionOrderByOperatorTest {
return (StreamingSelectionOrderByOperator) streamingOperator;
}
- /// Drains the streaming operator for {@code query} and returns all rows concatenated in emission order.
+ /// Drains the streaming operator for `query` and returns all rows concatenated in emission order.
private List<Object[]> collectStreamingRows(IndexSegment segment, @Language("sql") String query,
boolean nullHandling) {
QueryContext queryContext = QueryContextConverterUtils.getQueryContext(query);
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2Test.java b/pinot-core/src/test/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2Test.java
index 64c6126bc43..45270526ea1 100644
--- a/pinot-core/src/test/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2Test.java
+++ b/pinot-core/src/test/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2Test.java
@@ -29,11 +29,12 @@ import org.testng.annotations.Test;
import static org.testng.Assert.assertEquals;
-/// Tests for execution-thread resolution in [InstancePlanMakerImplV2], covering the interplay
-/// between `max.execution.threads`, `default.execution.threads`, and per-query overrides.
+/// Tests for server-config versus query-option resolution in [InstancePlanMakerImplV2]: execution threads
+/// (`max.execution.threads`, `default.execution.threads`) and the streaming selection ORDER BY merge defaults.
public class InstancePlanMakerImplV2Test {
private static final String BASE_QUERY = "SELECT * FROM testTable";
+ private static final String SELECTION_ORDER_BY_QUERY = "SELECT col FROM testTable ORDER BY col LIMIT 10";
private static QueryContext buildQueryContext(String maxExecutionThreadsOption) {
QueryContext queryContext = QueryContextConverterUtils.getQueryContext(BASE_QUERY);
@@ -99,4 +100,66 @@ public class InstancePlanMakerImplV2Test {
config.setProperty(Server.DEFAULT_EXECUTION_THREADS, 8);
planMaker.init(config);
}
+
+ @Test
+ public void testSortedSelectionMergeDefaultsWithoutServerConfig() {
+ InstancePlanMakerImplV2 planMaker = new InstancePlanMakerImplV2();
+ planMaker.init(new PinotConfiguration());
+
+ QueryContext queryContext = QueryContextConverterUtils.getQueryContext(SELECTION_ORDER_BY_QUERY);
+ planMaker.applyQueryOptions(queryContext);
+
+ assertEquals(queryContext.getSortedSelectionMergeAutoMinSortedRatio(),
+ Server.DEFAULT_SORTED_SELECTION_MERGE_AUTO_MIN_SORTED_RATIO);
+ assertEquals(queryContext.getSortedSelectionMergeBlockSize(), Server.DEFAULT_SORTED_SELECTION_MERGE_BLOCK_SIZE);
+ }
+
+ @Test
+ public void testSortedSelectionMergeServerConfigAppliesWithoutQueryOption() {
+ InstancePlanMakerImplV2 planMaker = new InstancePlanMakerImplV2();
+ planMaker.init(sortedSelectionMergeConfig(0.6, 2_500));
+
+ QueryContext queryContext = QueryContextConverterUtils.getQueryContext(SELECTION_ORDER_BY_QUERY);
+ planMaker.applyQueryOptions(queryContext);
+
+ assertEquals(queryContext.getSortedSelectionMergeAutoMinSortedRatio(), 0.6);
+ assertEquals(queryContext.getSortedSelectionMergeBlockSize(), 2_500);
+ }
+
+ @Test
+ public void testSortedSelectionMergeQueryOptionOverridesServerConfig() {
+ InstancePlanMakerImplV2 planMaker = new InstancePlanMakerImplV2();
+ planMaker.init(sortedSelectionMergeConfig(0.6, 2_500));
+
+ QueryContext queryContext = QueryContextConverterUtils.getQueryContext(SELECTION_ORDER_BY_QUERY);
+ Map<String, String> queryOptions = queryContext.getQueryOptions();
+ queryOptions.put(QueryOptionKey.SORTED_SELECTION_MERGE_AUTO_MIN_SORTED_RATIO, "0.9");
+ queryOptions.put(QueryOptionKey.SORTED_SELECTION_MERGE_BLOCK_SIZE, "500");
+ planMaker.applyQueryOptions(queryContext);
+
+ assertEquals(queryContext.getSortedSelectionMergeAutoMinSortedRatio(), 0.9);
+ assertEquals(queryContext.getSortedSelectionMergeBlockSize(), 500);
+ }
+
+ @Test(expectedExceptions = IllegalStateException.class)
+ public void testInitRejectsSortedSelectionMergeRatioAboveOne() {
+ new InstancePlanMakerImplV2().init(sortedSelectionMergeConfig(1.5, 2_500));
+ }
+
+ @Test(expectedExceptions = IllegalStateException.class)
+ public void testInitRejectsNegativeSortedSelectionMergeRatio() {
+ new InstancePlanMakerImplV2().init(sortedSelectionMergeConfig(-0.1, 2_500));
+ }
+
+ @Test(expectedExceptions = IllegalStateException.class)
+ public void testInitRejectsNonPositiveSortedSelectionMergeBlockSize() {
+ new InstancePlanMakerImplV2().init(sortedSelectionMergeConfig(0.6, 0));
+ }
+
+ private static PinotConfiguration sortedSelectionMergeConfig(double autoMinSortedRatio, int blockSize) {
+ PinotConfiguration config = new PinotConfiguration();
+ config.setProperty(Server.SORTED_SELECTION_MERGE_AUTO_MIN_SORTED_RATIO, autoMinSortedRatio);
+ config.setProperty(Server.SORTED_SELECTION_MERGE_BLOCK_SIZE, blockSize);
+ return config;
+ }
}
diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
index 4a8d90207b1..efbbbe15f1a 100644
--- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
+++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
@@ -631,10 +631,6 @@ public class CommonConstants {
"pinot.broker.mse.streaming.distinct.flush.threshold";
public static final int DEFAULT_MSE_STREAMING_DISTINCT_FLUSH_THRESHOLD = -1;
- /// Default output block size (rows) for the streaming selection ORDER BY combine
- /// ({@link Request.QueryOptionKey#SORTED_SELECTION_MERGE_BLOCK_SIZE}).
- public static final int DEFAULT_SORTED_SELECTION_MERGE_BLOCK_SIZE = 10_000;
-
// Whether to infer partition hint by default or not.
// This value can always be overridden by INFER_PARTITION_HINT query option
public static final String CONFIG_OF_INFER_PARTITION_HINT = "pinot.broker.multistage.infer.partition.hint";
@@ -1650,6 +1646,21 @@ public class CommonConstants {
/// among many sorted ones. This default is chosen to tolerate that while still rejecting a mostly-unsorted table.
/// It is a starting point rather than a measured optimum.
public static final double DEFAULT_SORTED_SELECTION_MERGE_AUTO_MIN_SORTED_RATIO = 0.8;
+ /// Server-wide override of [#DEFAULT_SORTED_SELECTION_MERGE_AUTO_MIN_SORTED_RATIO], applied to queries that do not
+ /// set the query option.
+ public static final String SORTED_SELECTION_MERGE_AUTO_MIN_SORTED_RATIO =
+ "sorted.selection.merge.auto.min.sorted.ratio";
+ public static final String CONFIG_OF_QUERY_EXECUTOR_SORTED_SELECTION_MERGE_AUTO_MIN_SORTED_RATIO =
+ QUERY_EXECUTOR_CONFIG_PREFIX + "." + SORTED_SELECTION_MERGE_AUTO_MIN_SORTED_RATIO;
+
+ /// Default output block size (rows) for the streaming selection ORDER BY combine
+ /// ([Broker.Request.QueryOptionKey#SORTED_SELECTION_MERGE_BLOCK_SIZE]).
+ public static final int DEFAULT_SORTED_SELECTION_MERGE_BLOCK_SIZE = 10_000;
+ /// Server-wide override of [#DEFAULT_SORTED_SELECTION_MERGE_BLOCK_SIZE], applied to queries that do not set the
+ /// query option.
+ public static final String SORTED_SELECTION_MERGE_BLOCK_SIZE = "sorted.selection.merge.block.size";
+ public static final String CONFIG_OF_QUERY_EXECUTOR_SORTED_SELECTION_MERGE_BLOCK_SIZE =
+ QUERY_EXECUTOR_CONFIG_PREFIX + "." + SORTED_SELECTION_MERGE_BLOCK_SIZE;
public static final String CONFIG_OF_MSE_MIN_GROUP_TRIM_SIZE = MSE_CONFIG_PREFIX + ".min.group.trim.size";
// Match the value of GroupByUtils.DEFAULT_MIN_NUM_GROUPS
public static final int DEFAULT_MSE_MIN_GROUP_TRIM_SIZE = 5000;
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
(pinot) 06/07: Check the DESC null-bearing case against fixed rows instead of MinMax
Posted by xi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
xiangfu0 pushed a commit to branch xiangfu0/codex1/sorted-exchange-leaf
in repository https://gitbox.apache.org/repos/asf/pinot.git
commit 164da0d2252cf57459d3741f906b4a54205bc4a6
Author: rohity <ro...@uber.com>
AuthorDate: Tue Sep 29 19:35:34 2026 +0000
Check the DESC null-bearing case against fixed rows instead of MinMax
MinMax prunes on the metadata max, which ignores nulls, so with fewer worker threads than segments it can skip the
null-bearing segment and drop the null rows that DESC puts first. That made the parity check fail on CI's smaller
runners. Assert the expected rows directly.
---
.../combine/StreamingSelectionOrderByCombineOperatorTest.java | 10 ++++++++--
1 file changed, 8 insertions(+), 2 deletions(-)
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java
index 2532d4aacc7..2950d4062bf 100644
--- a/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java
+++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java
@@ -21,6 +21,7 @@ package org.apache.pinot.core.operator.combine;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
+import java.util.Arrays;
import java.util.Comparator;
import java.util.List;
import java.util.Map;
@@ -732,11 +733,16 @@ public class StreamingSelectionOrderByCombineOperatorTest {
assertSegmentsScanned(result, 2, segments.size());
// The ASC count above does not prove this segment was force-activated: its stored null is Integer.MIN_VALUE, so
// the metadata min would activate it anyway. On DESC the metadata max is the highest non-null and sits below the
- // high range, so keeping that max would never acquire it. Its nulls sort first, which is what parity checks.
+ // high range, so keeping that max would never acquire it. Its nulls sort first on DESC (valCol DESC breaks ties),
+ // so they must lead the result. Checked against fixed rows, not MinMax: MinMax prunes on that same max, so with
+ // fewer worker threads than segments it can skip this segment and drop its nulls.
@Language("sql") String desc =
"SET allowReverseOrder=true; SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol DESC, valCol DESC "
+ "LIMIT 5";
- assertParity(segments, desc, true);
+ Result descResult = run(segments, desc, true, true, true, 3);
+ assertEquals(descResult._rows.stream().map(Arrays::asList).collect(Collectors.toList()),
+ List.of(Arrays.asList(null, 1002), Arrays.asList(null, 1001), Arrays.asList(null, 1000), List.of(2099, 2099),
+ List.of(2098, 2098)));
QueryContext descContext = hintedContext(desc);
descContext.setNullHandlingEnabled(true);
List<InstrumentedSegmentOperator> children = instrument(segments, descContext);
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
(pinot) 01/07: Add streaming selection ORDER BY combine for physically sorted segments
Posted by xi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
xiangfu0 pushed a commit to branch xiangfu0/codex1/sorted-exchange-leaf
in repository https://gitbox.apache.org/repos/asf/pinot.git
commit 849901a015daf0a479f8a281ca7b639e53e22078
Author: rohity <ro...@uber.com>
AuthorDate: Wed Aug 5 18:24:30 2026 +0000
Add streaming selection ORDER BY combine for physically sorted segments
An unbounded leaf-stage ORDER BY (as injected for sorted merge join inputs)
routes to MinMaxValueBasedSelectionOrderByCombineOperator, which merges every
segment's rows into a single block before returning anything. At large data
volumes this exceeds the leaf stage's CPU budget and ThreadAccountant raises
EarlyTerminationException inside SelectionOperatorUtils.mergeWithOrdering(),
surfacing to the broker as a spurious "Cancelled by sender".
This adds a streaming alternative, opt-in via the `streamingSelectionOrderBy`
query option:
- StreamingSelectionOrderByOperator emits sorted blocks incrementally for a
segment that is physically sorted on the leading ORDER BY column, reading
the sorted forward index in order instead of building a priority queue.
- StreamingSelectionOrderByCombineOperator performs a k-way heap merge across
segment operators and emits bounded blocks (`streamingSelectionOrderByBlockSize`,
default 10000) rather than one materialized result.
- SelectionPlanNode and CombinePlanNode select these operators when the option
is set and the sortedness precondition holds; otherwise behaviour is unchanged.
Part of #18667.
---
.../common/utils/config/QueryOptionsUtils.java | 10 +
.../common/utils/config/QueryOptionsUtilsTest.java | 4 +-
.../StreamingSelectionOrderByCombineOperator.java | 543 ++++++++++++++++++++
.../query/StreamingSelectionOrderByOperator.java | 514 +++++++++++++++++++
.../apache/pinot/core/plan/CombinePlanNode.java | 16 +
.../apache/pinot/core/plan/SelectionPlanNode.java | 28 +-
.../core/plan/maker/InstancePlanMakerImplV2.java | 11 +
.../core/query/request/context/QueryContext.java | 22 +
.../operator/combine/CombineSlowOperatorsTest.java | 94 +++-
...reamingSelectionOrderByCombineOperatorTest.java | 553 +++++++++++++++++++++
.../StreamingSelectionOrderByOperatorTest.java | 406 +++++++++++++++
.../apache/pinot/spi/utils/CommonConstants.java | 10 +
12 files changed, 2198 insertions(+), 13 deletions(-)
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java
index b028106ed68..d8393720613 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java
@@ -602,6 +602,16 @@ public class QueryOptionsUtils {
return checkedParseIntNonNegative(QueryOptionKey.STREAMING_DISTINCT_FLUSH_THRESHOLD, value);
}
+ public static boolean isSortedSelectionMergeEnabled(Map<String, String> queryOptions) {
+ return Boolean.parseBoolean(queryOptions.get(QueryOptionKey.SORTED_SELECTION_MERGE_ENABLED));
+ }
+
+ @Nullable
+ public static Integer getSortedSelectionMergeBlockSize(Map<String, String> queryOptions) {
+ String value = queryOptions.get(QueryOptionKey.SORTED_SELECTION_MERGE_BLOCK_SIZE);
+ return checkedParseIntPositive(QueryOptionKey.SORTED_SELECTION_MERGE_BLOCK_SIZE, value);
+ }
+
public static boolean isNullHandlingEnabled(Map<String, String> queryOptions) {
return Boolean.parseBoolean(queryOptions.get(QueryOptionKey.ENABLE_NULL_HANDLING));
}
diff --git a/pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java b/pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java
index f2125adef28..ffb698e789d 100644
--- a/pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java
+++ b/pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java
@@ -39,7 +39,7 @@ import static org.testng.Assert.fail;
public class QueryOptionsUtilsTest {
private static final List<String> POSITIVE_INT_KEYS =
List.of(NUM_REPLICA_GROUPS_TO_QUERY, MAX_EXECUTION_THREADS, NUM_GROUPS_LIMIT, MAX_INITIAL_RESULT_HOLDER_CAPACITY,
- MAX_STREAMING_PENDING_BLOCKS, MAX_ROWS_IN_JOIN, MAX_ROWS_IN_WINDOW);
+ MAX_STREAMING_PENDING_BLOCKS, MAX_ROWS_IN_JOIN, MAX_ROWS_IN_WINDOW, SORTED_SELECTION_MERGE_BLOCK_SIZE);
private static final List<String> NON_NEGATIVE_INT_KEYS =
List.of(MULTI_STAGE_LEAF_LIMIT, STREAMING_GROUP_BY_FLUSH_THRESHOLD, STREAMING_DISTINCT_FLUSH_THRESHOLD);
private static final List<String> UNBOUNDED_INT_KEYS =
@@ -375,6 +375,8 @@ public class QueryOptionsUtilsTest {
return QueryOptionsUtils.getMaxRowsInJoin(map);
case MAX_ROWS_IN_WINDOW:
return QueryOptionsUtils.getMaxRowsInWindow(map);
+ case SORTED_SELECTION_MERGE_BLOCK_SIZE:
+ return QueryOptionsUtils.getSortedSelectionMergeBlockSize(map);
// Non-negative ints
case MULTI_STAGE_LEAF_LIMIT:
return QueryOptionsUtils.getMultiStageLeafLimit(map);
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperator.java
new file mode 100644
index 00000000000..b99d6f0edc7
--- /dev/null
+++ b/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();
+ if (cursor == null) {
+ // All active cursors exhausted (and pruning guarantees the rest cannot contribute).
+ break;
+ }
+ _outputRows.add(cursor.currentHead());
+ _numRowsEmitted++;
+ cursor.advance();
+ if (cursor.currentHead() != null) {
+ _priorityQueue.offer(cursor);
+ }
+ if (_streaming && _outputRows.size() >= _blockSize) {
+ return flushDataBlock();
+ }
+ }
+ // Merge complete.
+ finish();
+ if (!_outputRows.isEmpty()) {
+ // Streaming: the final partial data block (next call returns the terminal metadata block).
+ // Single-stage: the single complete block.
+ return flushDataBlock();
+ }
+ if (_streaming) {
+ return attachDataSchemaMismatchErrors(attachExecutionStats(new MetadataResultsBlock()));
+ }
+ // Single-stage with no rows: still return a (single) block carrying the schema and execution stats.
+ return attachDataSchemaMismatchErrors(attachExecutionStats(
+ new SelectionResultsBlock(resolveDataSchema(), List.of(), _comparator, _queryContext)));
+ } catch (Exception e) {
+ _done = true;
+ releaseAllCursors();
+ return createExceptionResultsBlockAndAttachExecutionStats(e, "merging sorted selection results");
+ } catch (Throwable t) {
+ // An Error (e.g. OutOfMemoryError while accumulating output rows) must not leave segments acquired for the
+ // lifetime of the server. In single-stage mode there is no start()/stop() backstop around this operator.
+ _done = true;
+ releaseAllCursors();
+ throw t;
+ }
+ }
+
+ /// Activates not-yet-active cursors whose min/max value can still contribute before the current merge frontier.
+ ///
+ /// Cursors are visited in min/max-sorted order and {@code _nextToActivate} is advanced only when a cursor is
+ /// actually activated; a {@code break} merely defers the current cursor, which is re-evaluated against the (rising)
+ /// frontier on every subsequent call. Correctness: a not-yet-activated cursor whose first order-by value range starts
+ /// strictly beyond the current heap head cannot contain a row that sorts before that head (the first order-by column
+ /// is the primary sort key), so the head is the true global minimum and is safe to emit; the deferred cursor is
+ /// activated later, exactly when the frontier reaches its min/max. When the heap is empty the frontier is unknown, so
+ /// activation is forced (never pruned), which also drains any segments that sort entirely after the ones seen so far.
+ private void activateEligibleCursors() {
+ while (_nextToActivate < _sortedCursors.length) {
+ SegmentCursor cursor = _sortedCursors[_nextToActivate];
+ if (_pruningEnabled) {
+ Comparable bound = _asc ? cursor._minValue : cursor._maxValue;
+ // A null bound means the segment must always be processed. Otherwise, only prune against a non-null frontier;
+ // if the head's first order-by value is null we cannot compare, so fall through and activate.
+ if (bound != null && !_priorityQueue.isEmpty()) {
+ Object headValue = _priorityQueue.peek().currentHead()[0];
+ if (headValue != null) {
+ // Both come from the same first order-by column: the metadata min/max and the materialized row[0] share
+ // the column's stored type, so this comparison is type-safe (same assumption as MinMaxValueBased...).
+ int cmp = bound.compareTo(headValue);
+ if (_asc ? cmp > 0 : cmp < 0) {
+ break;
+ }
+ }
+ }
+ }
+ cursor.activate();
+ _nextToActivate++;
+ if (cursor.currentHead() != null) {
+ _priorityQueue.offer(cursor);
+ }
+ }
+ }
+
+ /// Marks the merge done and releases any segments still held by un-drained cursors (e.g. when LIMIT is reached).
+ private void finish() {
+ _done = true;
+ releaseAllCursors();
+ }
+
+ private void releaseAllCursors() {
+ for (SegmentCursor cursor : _sortedCursors) {
+ cursor.release();
+ }
+ }
+
+ /// Returns the accumulated output rows as a sorted {@link SelectionResultsBlock} and resets the output buffer. The
+ /// block carries the comparator so the broker-side n-way reduce stays correct. In single-stage mode it is the only
+ /// block, so execution stats are attached; in streaming mode stats are attached to the terminal metadata block.
+ private BaseResultsBlock flushDataBlock() {
+ List<Object[]> rows = _outputRows;
+ _outputRows = newOutputList();
+ SelectionResultsBlock block = new SelectionResultsBlock(resolveDataSchema(), rows, _comparator, _queryContext);
+ attachDataSchemaMismatchErrors(block);
+ return _streaming ? block : attachExecutionStats(block);
+ }
+
+ /// Records that a segment's block was dropped because its schema disagreed with the merge schema. Deduplicated by
+ /// message so a mid-reload table with many divergent segments cannot flood the response.
+ private void recordDataSchemaMismatch(@Nullable DataSchema mismatched) {
+ String errorMessage =
+ String.format("Data schema mismatch between merged block: %s and block to merge: %s, drop block to merge",
+ _dataSchema, mismatched);
+ // NOTE: This is segment level log, so log at debug level to prevent flooding the log.
+ LOGGER.debug(errorMessage);
+ if (_dataSchemaMismatchErrors == null) {
+ _dataSchemaMismatchErrors = new HashSet<>();
+ _unreportedDataSchemaMismatchErrors = new ArrayList<>();
+ }
+ if (_dataSchemaMismatchErrors.add(errorMessage)) {
+ _unreportedDataSchemaMismatchErrors.add(errorMessage);
+ }
+ }
+
+ /// Attaches mismatch errors recorded since the last emitted block. In streaming mode blocks are emitted many times,
+ /// so errors are drained rather than re-attached, and the terminal block picks up any recorded after the last flush.
+ private <T extends BaseResultsBlock> T attachDataSchemaMismatchErrors(T block) {
+ if (_unreportedDataSchemaMismatchErrors != null && !_unreportedDataSchemaMismatchErrors.isEmpty()) {
+ for (String errorMessage : _unreportedDataSchemaMismatchErrors) {
+ block.addErrorMessage(QueryErrorMessage.safeMsg(QueryErrorCode.MERGE_RESPONSE, errorMessage));
+ }
+ _unreportedDataSchemaMismatchErrors.clear();
+ }
+ return block;
+ }
+
+ private List<Object[]> newOutputList() {
+ int capacity =
+ Math.min(_blockSize, Math.min(_numRowsToKeep, SelectionOperatorUtils.MAX_ROW_HOLDER_INITIAL_CAPACITY));
+ return new ArrayList<>(Math.max(1, capacity));
+ }
+
+ /// Returns the data schema captured from the first child block. If no segment produced a block (every segment is a
+ /// streaming operator that matched zero rows), reconstructs the schema from the first segment so that an empty result
+ /// still carries a valid, correctly-ordered schema (order-by expressions first, matching the child blocks' layout).
+ private DataSchema resolveDataSchema() {
+ if (_dataSchema != null) {
+ return _dataSchema;
+ }
+ IndexSegment indexSegment = _operators.get(0).getIndexSegment();
+ List<ExpressionContext> expressions = SelectionOperatorUtils.extractExpressions(_queryContext, indexSegment);
+ Set<String> columns = new HashSet<>();
+ for (ExpressionContext expression : expressions) {
+ expression.getColumns(columns);
+ }
+ Map<String, DataSource> dataSourceMap = new HashMap<>();
+ for (String column : columns) {
+ dataSourceMap.put(column, indexSegment.getDataSource(column, _queryContext.getSchema()));
+ }
+ int numExpressions = expressions.size();
+ String[] columnNames = new String[numExpressions];
+ ColumnDataType[] columnDataTypes = new ColumnDataType[numExpressions];
+ for (int i = 0; i < numExpressions; i++) {
+ ExpressionContext expression = expressions.get(i);
+ columnNames[i] = expression.toString();
+ TransformFunction transformFunction = TransformFunctionFactory.get(expression, dataSourceMap);
+ columnDataTypes[i] = ColumnDataType.fromDataType(transformFunction.getResultMetadata().getDataType(),
+ transformFunction.getResultMetadata().isSingleValue());
+ }
+ _dataSchema = new DataSchema(columnNames, columnDataTypes);
+ return _dataSchema;
+ }
+
+ /// Iterates a single segment's locally-sorted rows, pulling blocks lazily from its operator. Streaming-backed cursors
+ /// loop until the operator returns {@code null}; single-block-backed cursors read exactly one block. The segment is
+ /// acquired on activation and released once exhausted or when the combine finishes (see
+ /// {@link StreamingSelectionOrderByCombineOperator}).
+ private class SegmentCursor {
+ private final Operator<BaseResultsBlock> _operator;
+ @Nullable
+ private final Comparable _minValue;
+ @Nullable
+ private final Comparable _maxValue;
+
+ private boolean _streamingChild;
+ private List<Object[]> _rows;
+ private int _pos;
+ private boolean _acquired;
+ /// Cached current head (rows.get(pos)) so the hot heap comparator does not re-index per comparison
+ @Nullable
+ private Object[] _head;
+
+ SegmentCursor(Operator<BaseResultsBlock> operator, @Nullable Comparable minValue, @Nullable Comparable maxValue) {
+ _operator = operator;
+ _minValue = minValue;
+ _maxValue = maxValue;
+ }
+
+ /// Returns the current head row to be merged next, or {@code null} if not activated or exhausted.
+ @Nullable
+ Object[] currentHead() {
+ return _head;
+ }
+
+ /// Acquires the segment, resolves whether the child is the lazy streaming operator, and reads the first block.
+ /// After this call {@link #currentHead()} returns the first row, or {@code null} if the segment contributes
+ /// nothing.
+ void activate() {
+ acquireSegment();
+ _streamingChild = isStreamingChild();
+ if (!pullBlock()) {
+ exhaust();
+ }
+ }
+
+ /// Advances past the current head, pulling the next block lazily for streaming cursors. Releases the segment
+ /// when the cursor is exhausted; afterwards {@link #currentHead()} returns {@code null}.
+ void advance() {
+ _pos++;
+ if (_pos < _rows.size()) {
+ _head = _rows.get(_pos);
+ return;
+ }
+ // Current block drained: streaming cursors pull the next run/block; single-block cursors are done.
+ if (_streamingChild && pullBlock()) {
+ return;
+ }
+ exhaust();
+ }
+
+ /// Loads the next non-empty block of rows from the operator, capturing the combine-level data schema on first
+ /// sight. Returns {@code false} when the operator is exhausted (no more rows). Streaming-backed operators emit
+ /// one run/block per call and {@code null} when done; single-block operators emit a single block and must not be
+ /// called again afterwards, so an empty/null block from a single-block child is treated as exhausted.
+ ///
+ /// A block whose schema differs from the one already captured is dropped and reported, mirroring
+ /// {@link org.apache.pinot.core.operator.combine.merger.SelectionOrderByResultsBlockMerger}. Segments on a server
+ /// can disagree on schema mid-reload (a newly added column exists only in reloaded segments), and merging rows of
+ /// differing width under one schema would corrupt the result rather than fail.
+ private boolean pullBlock() {
+ while (true) {
+ SelectionResultsBlock block = nextBlock();
+ if (block == null) {
+ return false;
+ }
+ if (_dataSchema == null) {
+ _dataSchema = block.getDataSchema();
+ } else if (!_dataSchema.equals(block.getDataSchema())) {
+ recordDataSchemaMismatch(block.getDataSchema());
+ return false;
+ }
+ List<Object[]> rows = block.getRows();
+ if (rows != null && !rows.isEmpty()) {
+ _rows = rows;
+ _pos = 0;
+ _head = rows.get(0);
+ return true;
+ }
+ // Defensive: an unexpected empty (non-null) block. Keep pulling only for streaming children; a single-block
+ // child yields exactly one block, so treat it as exhausted.
+ if (!_streamingChild) {
+ return false;
+ }
+ }
+ }
+
+ private SelectionResultsBlock nextBlock() {
+ try {
+ return (SelectionResultsBlock) _operator.nextBlock();
+ } catch (RuntimeException e) {
+ throw wrapOperatorException(_operator, e);
+ }
+ }
+
+ /// Returns whether the underlying child operator is the lazy {@link StreamingSelectionOrderByOperator}. Must be
+ /// called after {@link #acquireSegment()} because materializing the wrapped child runs the plan node, which
+ /// accesses segment buffers.
+ private boolean isStreamingChild() {
+ Operator underlying = _operator;
+ if (_operator instanceof AcquireReleaseColumnsSegmentOperator) {
+ AcquireReleaseColumnsSegmentOperator wrapper = (AcquireReleaseColumnsSegmentOperator) _operator;
+ wrapper.materializeChildOperator();
+ underlying = wrapper.getChildOperators().get(0);
+ }
+ return underlying instanceof StreamingSelectionOrderByOperator;
+ }
+
+ private void acquireSegment() {
+ if (_operator instanceof AcquireReleaseColumnsSegmentOperator) {
+ ((AcquireReleaseColumnsSegmentOperator) _operator).acquire();
+ }
+ _acquired = true;
+ }
+
+ /// Releases the segment if still held. Idempotent: safe to call from {@link #exhaust()} and combine cleanup.
+ private void release() {
+ if (_acquired) {
+ if (_operator instanceof AcquireReleaseColumnsSegmentOperator) {
+ ((AcquireReleaseColumnsSegmentOperator) _operator).release();
+ }
+ _acquired = false;
+ }
+ }
+
+ private void exhaust() {
+ release();
+ _rows = null;
+ _head = null;
+ }
+ }
+}
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperator.java
new file mode 100644
index 00000000000..e66a6ee0391
--- /dev/null
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperator.java
@@ -0,0 +1,514 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.core.operator.query;
+
+import com.google.common.base.CaseFormat;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.PriorityQueue;
+import java.util.Set;
+import java.util.stream.Collectors;
+import javax.annotation.Nullable;
+import org.apache.pinot.common.request.context.ExpressionContext;
+import org.apache.pinot.common.request.context.OrderByExpressionContext;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.core.common.BlockValSet;
+import org.apache.pinot.core.common.Operator;
+import org.apache.pinot.core.common.RowBasedBlockValueFetcher;
+import org.apache.pinot.core.operator.BaseOperator;
+import org.apache.pinot.core.operator.BaseProjectOperator;
+import org.apache.pinot.core.operator.BitmapDocIdSetOperator;
+import org.apache.pinot.core.operator.ColumnContext;
+import org.apache.pinot.core.operator.ExecutionStatistics;
+import org.apache.pinot.core.operator.ExplainAttributeBuilder;
+import org.apache.pinot.core.operator.ProjectionOperator;
+import org.apache.pinot.core.operator.ProjectionOperatorUtils;
+import org.apache.pinot.core.operator.blocks.ValueBlock;
+import org.apache.pinot.core.operator.blocks.results.SelectionResultsBlock;
+import org.apache.pinot.core.operator.transform.TransformOperator;
+import org.apache.pinot.core.query.request.context.QueryContext;
+import org.apache.pinot.core.query.selection.SelectionOperatorUtils;
+import org.apache.pinot.core.query.utils.OrderByComparatorFactory;
+import org.apache.pinot.segment.spi.IndexSegment;
+import org.apache.pinot.segment.spi.datasource.DataSource;
+import org.apache.pinot.spi.query.QueryScanCostContext;
+import org.roaringbitmap.RoaringBitmap;
+
+
+/// Lazy, incremental selection ORDER BY operator for segments that are physically sorted on the first order-by column.
+///
+/// Unlike {@link SelectionOrderByOperator} (which materializes the segment's whole top-K in a single block) this
+/// operator emits one globally-sorted {@link SelectionResultsBlock} per {@link #getNextBlock()} call and returns
+/// {@code null} when the segment is exhausted, so that a downstream k-way-merge combine operator can pull from many
+/// segments lazily and stop early. It relies on the underlying project operator iterating the first order-by column in
+/// the query order (the caller must guarantee {@code projectOperator.isCompatibleWith(DocIdOrder.fromAsc(asc))}).
+///
+/// It runs in one of two emission modes:
+///
+/// - **No tail to sort** ({@code numSortedExpressions == numOrderByExpressions}, e.g. {@code ORDER BY sorted}):
+/// rows already arrive from the project operator in final order, so each call emits the next project block
+/// (trimmed to the remaining {@code limit + offset} budget).
+/// - **Tail to sort** ({@code numSortedExpressions < numOrderByExpressions}, e.g.
+/// {@code ORDER BY sorted, other}):
+/// each call reads forward until the first order-by value changes (a primary-value "run"), retains the run's top
+/// {@code limit + offset} rows by the full comparator, and emits them sorted. This bounds the in-memory run buffer to
+/// {@code limit + offset} rows even when the first order-by column is near-constant (very low cardinality).
+///
+/// Like {@link SelectionOrderByOperator} it preserves the two-phase projection optimization: when there are output
+/// expressions that are not order-by expressions, the forward scan only fetches the order-by expressions plus the
+/// document id, and the non-order-by expressions are fetched in a second pass over the retained document ids of each
+/// emitted block.
+///
+/// This operator is stateful across {@link #getNextBlock()} calls and is **not** thread-safe; a single consumer
+/// must drive it.
+public class StreamingSelectionOrderByOperator extends BaseOperator<SelectionResultsBlock> {
+ private static final String EXPLAIN_NAME = "SELECT_ORDERBY_STREAMING";
+
+ private final IndexSegment _indexSegment;
+ private final QueryContext _queryContext;
+ private final boolean _nullHandlingEnabled;
+ /// Deduped order-by expressions followed by output expressions from SelectionOperatorUtils.extractExpressions()
+ private final List<ExpressionContext> _expressions;
+ private final BaseProjectOperator<?> _projectOperator;
+ private final List<OrderByExpressionContext> _orderByExpressions;
+ private final ColumnContext[] _orderByColumnContexts;
+ private final int _numExpressions;
+ private final int _numOrderByExpressions;
+ private final int _numRowsToKeep;
+ /// Whether there are output expressions that are not order-by expressions (requires the two-phase fetch)
+ private final boolean _twoPhase;
+ /// Whether the order-by has an unsorted tail that must be sorted in memory per run
+ private final boolean _tailToSort;
+ /// Expressions fetched during the forward scan: order-by expressions only when two-phase, otherwise all expressions
+ private final List<ExpressionContext> _phase1Expressions;
+ private final int _numPhase1Columns;
+ private final Comparator<Object[]> _comparator;
+ /// Compares only the first order-by column; used to detect primary-value run boundaries
+ private final Comparator<Object[]> _primaryComparator;
+ /// Pre-allocated run heap (cleared and reused each nextRun() call to avoid per-run allocation)
+ private final Comparator<Object[]> _reversedComparator;
+ private final PriorityQueue<Object[]> _runHeap;
+
+ // Pre-computed invariants for the two-phase fetch (null when single-phase)
+ private final List<ExpressionContext> _nonOrderByExpressions;
+ private final Map<String, DataSource> _phase2DataSourceMap;
+ private final int _phase2NumColumns;
+
+ /// Lazily built and cached; for two-phase it requires the transform operator's result column contexts
+ private DataSchema _dataSchema;
+
+ // Forward-scan cursor state (used by the tail-to-sort mode)
+ private ValueBlock _currentBlock;
+ private RowBasedBlockValueFetcher _currentFetcher;
+ private int[] _currentDocIds;
+ private RoaringBitmap[] _currentNullBitmaps;
+ private int _currentNumDocs;
+ private int _currentPos;
+ /// One-row lookahead: the first row of the next run, stashed when a run boundary is crossed
+ private Object[] _pendingRow;
+ private boolean _projectExhausted;
+
+ private boolean _exhausted;
+ private int _numRowsEmitted;
+ private int _numDocsScanned = 0;
+ private long _numEntriesScannedPostFilter = 0;
+
+ public StreamingSelectionOrderByOperator(IndexSegment indexSegment, QueryContext queryContext,
+ List<ExpressionContext> expressions, BaseProjectOperator<?> projectOperator, int numSortedExpressions) {
+ _indexSegment = indexSegment;
+ _queryContext = queryContext;
+ _nullHandlingEnabled = queryContext.isNullHandlingEnabled();
+ _expressions = expressions;
+ _projectOperator = projectOperator;
+
+ _orderByExpressions = queryContext.getOrderByExpressions();
+ assert _orderByExpressions != null;
+ _numExpressions = expressions.size();
+ _numOrderByExpressions = _orderByExpressions.size();
+ _orderByColumnContexts = new ColumnContext[_numOrderByExpressions];
+ for (int i = 0; i < _numOrderByExpressions; i++) {
+ ExpressionContext expression = _orderByExpressions.get(i).getExpression();
+ _orderByColumnContexts[i] = _projectOperator.getResultColumnContext(expression);
+ }
+
+ _numRowsToKeep = queryContext.getOffset() + queryContext.getLimit();
+ _twoPhase = _numExpressions > _numOrderByExpressions;
+ _tailToSort = numSortedExpressions < _numOrderByExpressions;
+ _comparator =
+ OrderByComparatorFactory.getComparator(_orderByExpressions, _orderByColumnContexts, _nullHandlingEnabled);
+ // The first order-by column is the physically sorted column, so it never contains nulls on this path; comparing
+ // only index 0 is enough to detect when one primary-value run ends and the next begins.
+ _primaryComparator =
+ OrderByComparatorFactory.getComparator(_orderByExpressions, _orderByColumnContexts, _nullHandlingEnabled, 0, 1);
+ _reversedComparator = _comparator.reversed();
+ _runHeap = new PriorityQueue<>(
+ Math.min(_numRowsToKeep, SelectionOperatorUtils.MAX_ROW_HOLDER_INITIAL_CAPACITY), _reversedComparator);
+
+ if (_twoPhase) {
+ _phase1Expressions = new ArrayList<>(_numOrderByExpressions);
+ for (OrderByExpressionContext orderByExpression : _orderByExpressions) {
+ _phase1Expressions.add(orderByExpression.getExpression());
+ }
+ _nonOrderByExpressions = _expressions.subList(_numOrderByExpressions, _numExpressions);
+ Set<String> columns = new HashSet<>();
+ for (ExpressionContext expressionContext : _nonOrderByExpressions) {
+ expressionContext.getColumns(columns);
+ }
+ _phase2NumColumns = columns.size();
+ _phase2DataSourceMap = new HashMap<>();
+ for (String column : columns) {
+ _phase2DataSourceMap.put(column, _indexSegment.getDataSource(column, _queryContext.getSchema()));
+ }
+ } else {
+ _phase1Expressions = _expressions;
+ _nonOrderByExpressions = null;
+ _phase2NumColumns = 0;
+ _phase2DataSourceMap = null;
+ // Single-phase: all output expressions are order-by expressions, so their types are known up front.
+ _dataSchema = buildSinglePhaseDataSchema();
+ }
+ _numPhase1Columns = _phase1Expressions.size();
+ }
+
+ @Override
+ protected SelectionResultsBlock getNextBlock() {
+ if (_exhausted) {
+ return null;
+ }
+ List<Object[]> rows = _tailToSort ? nextRun() : nextSortedRows();
+ if (rows == null || rows.isEmpty()) {
+ _exhausted = true;
+ return null;
+ }
+ if (_twoPhase) {
+ fetchNonOrderByColumns(rows);
+ }
+ // Single-phase builds the schema in the constructor; two-phase builds it during fetchNonOrderByColumns above.
+ assert _dataSchema != null;
+ return new SelectionResultsBlock(_dataSchema, rows, _comparator, _queryContext);
+ }
+
+ /// No-tail-to-sort mode: the project operator already returns rows in final order, so emit the next project block,
+ /// trimmed to the remaining {@code limit + offset} budget. Returns {@code null} when exhausted.
+ @Nullable
+ private List<Object[]> nextSortedRows() {
+ int remaining = _numRowsToKeep - _numRowsEmitted;
+ if (remaining <= 0) {
+ return null;
+ }
+ ValueBlock valueBlock = _projectOperator.nextBlock();
+ if (valueBlock == null) {
+ return null;
+ }
+ int numDocsFetched = valueBlock.getNumDocs();
+ BlockValSet[] blockValSets = new BlockValSet[_numPhase1Columns];
+ for (int i = 0; i < _numPhase1Columns; i++) {
+ blockValSets[i] = valueBlock.getBlockValueSet(_phase1Expressions.get(i));
+ }
+ RowBasedBlockValueFetcher blockValueFetcher = new RowBasedBlockValueFetcher(blockValSets);
+ int[] docIds = _twoPhase ? valueBlock.getDocIds() : null;
+ RoaringBitmap[] nullBitmaps = null;
+ if (_nullHandlingEnabled) {
+ nullBitmaps = new RoaringBitmap[_numPhase1Columns];
+ for (int i = 0; i < _numPhase1Columns; i++) {
+ nullBitmaps[i] = blockValSets[i].getNullBitmap();
+ }
+ }
+ _numDocsScanned += numDocsFetched;
+ _numEntriesScannedPostFilter += (long) numDocsFetched * _projectOperator.getNumColumnsProjected();
+ reportScanCost(numDocsFetched, (long) numDocsFetched * _projectOperator.getNumColumnsProjected());
+
+ // Rows arrive sorted; we only need the first 'remaining' of them globally.
+ int numRows = Math.min(numDocsFetched, remaining);
+ List<Object[]> rows = new ArrayList<>(numRows);
+ for (int i = 0; i < numRows; i++) {
+ rows.add(materializeRow(blockValueFetcher, docIds, nullBitmaps, i));
+ }
+ _numRowsEmitted += rows.size();
+ return rows;
+ }
+
+ /// Tail-to-sort mode: read forward until the first order-by value changes, retain the run's top
+ /// {@code limit + offset} rows by the full comparator, and return them sorted. Returns {@code null} when
+ /// exhausted.
+ @Nullable
+ private List<Object[]> nextRun() {
+ int remaining = _numRowsToKeep - _numRowsEmitted;
+ if (remaining <= 0) {
+ return null;
+ }
+ if (_pendingRow == null) {
+ _pendingRow = nextRow();
+ if (_pendingRow == null) {
+ return null;
+ }
+ }
+ PriorityQueue<Object[]> runHeap = _runHeap;
+ runHeap.clear();
+ Object[] runFirstRow = _pendingRow;
+ SelectionOperatorUtils.addToPriorityQueue(_pendingRow, runHeap, _numRowsToKeep);
+ _pendingRow = null;
+ Object[] row;
+ while ((row = nextRow()) != null) {
+ if (_primaryComparator.compare(row, runFirstRow) == 0) {
+ SelectionOperatorUtils.addToPriorityQueue(row, runHeap, _numRowsToKeep);
+ } else {
+ // Run boundary: this row starts the next run, keep it for the next call.
+ _pendingRow = row;
+ break;
+ }
+ }
+ List<Object[]> rows = drainAscending(runHeap);
+ // A segment never contributes more than 'limit + offset' rows to the global result, and they are a prefix of its
+ // local sorted order, so cap the total emitted across runs at the remaining budget (the rows are ascending, keep
+ // the smallest 'remaining').
+ if (rows.size() > remaining) {
+ rows = rows.subList(0, remaining);
+ }
+ _numRowsEmitted += rows.size();
+ return rows;
+ }
+
+ /// Pulls the next row of the forward scan (across project blocks), materialized as an
+ /// {@code Object[_numExpressions]}. For two-phase the document id is stashed at index
+ /// {@code _numOrderByExpressions} (overwritten in the second pass). Returns {@code null} when the project operator
+ /// is exhausted.
+ @Nullable
+ private Object[] nextRow() {
+ while (true) {
+ if (_currentBlock == null || _currentPos >= _currentNumDocs) {
+ if (_projectExhausted) {
+ return null;
+ }
+ _currentBlock = _projectOperator.nextBlock();
+ if (_currentBlock == null) {
+ _projectExhausted = true;
+ return null;
+ }
+ BlockValSet[] blockValSets = new BlockValSet[_numPhase1Columns];
+ for (int i = 0; i < _numPhase1Columns; i++) {
+ blockValSets[i] = _currentBlock.getBlockValueSet(_phase1Expressions.get(i));
+ }
+ _currentFetcher = new RowBasedBlockValueFetcher(blockValSets);
+ _currentNumDocs = _currentBlock.getNumDocs();
+ _currentDocIds = _twoPhase ? _currentBlock.getDocIds() : null;
+ if (_nullHandlingEnabled) {
+ _currentNullBitmaps = new RoaringBitmap[_numPhase1Columns];
+ for (int i = 0; i < _numPhase1Columns; i++) {
+ _currentNullBitmaps[i] = blockValSets[i].getNullBitmap();
+ }
+ }
+ _currentPos = 0;
+ _numDocsScanned += _currentNumDocs;
+ _numEntriesScannedPostFilter += (long) _currentNumDocs * _projectOperator.getNumColumnsProjected();
+ reportScanCost(_currentNumDocs, (long) _currentNumDocs * _projectOperator.getNumColumnsProjected());
+ if (_currentNumDocs == 0) {
+ _currentBlock = null;
+ continue;
+ }
+ }
+ int rowId = _currentPos++;
+ return materializeRow(_currentFetcher, _currentDocIds, _currentNullBitmaps, rowId);
+ }
+ }
+
+ /// Materializes a single phase-1 row (deep-copied out of the value block buffers) from the given fetcher.
+ private Object[] materializeRow(RowBasedBlockValueFetcher fetcher, @Nullable int[] docIds,
+ @Nullable RoaringBitmap[] nullBitmaps, int rowId) {
+ Object[] row = new Object[_numExpressions];
+ fetcher.getRow(rowId, row, 0);
+ if (_twoPhase) {
+ row[_numOrderByExpressions] = docIds[rowId];
+ }
+ if (_nullHandlingEnabled) {
+ for (int colId = 0; colId < _numPhase1Columns; colId++) {
+ if (nullBitmaps[colId] != null && nullBitmaps[colId].contains(rowId)) {
+ row[colId] = null;
+ }
+ }
+ }
+ return row;
+ }
+
+ /// Drains a max-heap (created with the reversed comparator) into an ascending list, mutable so the second pass can
+ /// fill non-order-by values in place.
+ private List<Object[]> drainAscending(PriorityQueue<Object[]> heap) {
+ int numRows = heap.size();
+ Object[][] sortedRows = new Object[numRows][];
+ for (int i = numRows - 1; i >= 0; i--) {
+ sortedRows[i] = heap.poll();
+ }
+ return Arrays.asList(sortedRows);
+ }
+
+ /// 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 {@link 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);
+
+ int numNonOrderByExpressions = _nonOrderByExpressions.size();
+ BlockValSet[] blockValSets = new BlockValSet[numNonOrderByExpressions];
+ int rowBaseId = 0;
+ ValueBlock valueBlock;
+ while ((valueBlock = transformOperator.nextBlock()) != null) {
+ for (int i = 0; i < numNonOrderByExpressions; i++) {
+ blockValSets[i] = valueBlock.getBlockValueSet(_nonOrderByExpressions.get(i));
+ }
+ RowBasedBlockValueFetcher blockValueFetcher = new RowBasedBlockValueFetcher(blockValSets);
+ int numDocsFetched = valueBlock.getNumDocs();
+ for (int i = 0; i < numDocsFetched; i++) {
+ blockValueFetcher.getRow(i, rowsByDocId.get(rowBaseId + i), _numOrderByExpressions);
+ }
+ if (_nullHandlingEnabled) {
+ RoaringBitmap[] nullBitmaps = new RoaringBitmap[numNonOrderByExpressions];
+ for (int i = 0; i < numNonOrderByExpressions; i++) {
+ nullBitmaps[i] = blockValSets[i].getNullBitmap();
+ }
+ for (int i = 0; i < numDocsFetched; i++) {
+ Object[] values = rowsByDocId.get(rowBaseId + i);
+ for (int colId = 0; colId < numNonOrderByExpressions; colId++) {
+ if (nullBitmaps[colId] != null && nullBitmaps[colId].contains(i)) {
+ values[_numOrderByExpressions + colId] = null;
+ }
+ }
+ }
+ }
+ _numEntriesScannedPostFilter += (long) numDocsFetched * _phase2NumColumns;
+ // Phase 2 re-reads docs already counted in phase 1, so only the extra entries are reported.
+ reportScanCost(0, (long) numDocsFetched * _phase2NumColumns);
+ rowBaseId += numDocsFetched;
+ }
+
+ if (_dataSchema == null) {
+ _dataSchema = buildTwoPhaseDataSchema(transformOperator);
+ }
+ }
+ }
+
+ private DataSchema buildSinglePhaseDataSchema() {
+ String[] columnNames = new String[_numExpressions];
+ DataSchema.ColumnDataType[] columnDataTypes = new DataSchema.ColumnDataType[_numExpressions];
+ for (int i = 0; i < _numExpressions; i++) {
+ columnNames[i] = _expressions.get(i).toString();
+ columnDataTypes[i] = DataSchema.ColumnDataType.fromDataType(_orderByColumnContexts[i].getDataType(),
+ _orderByColumnContexts[i].isSingleValue());
+ }
+ return new DataSchema(columnNames, columnDataTypes);
+ }
+
+ private DataSchema buildTwoPhaseDataSchema(TransformOperator transformOperator) {
+ int numNonOrderByExpressions = _nonOrderByExpressions.size();
+ String[] columnNames = new String[_numExpressions];
+ DataSchema.ColumnDataType[] columnDataTypes = new DataSchema.ColumnDataType[_numExpressions];
+ for (int i = 0; i < _numExpressions; i++) {
+ columnNames[i] = _expressions.get(i).toString();
+ }
+ for (int i = 0; i < _numOrderByExpressions; i++) {
+ columnDataTypes[i] = DataSchema.ColumnDataType.fromDataType(_orderByColumnContexts[i].getDataType(),
+ _orderByColumnContexts[i].isSingleValue());
+ }
+ for (int i = 0; i < numNonOrderByExpressions; i++) {
+ ColumnContext columnContext = transformOperator.getResultColumnContext(_nonOrderByExpressions.get(i));
+ columnDataTypes[_numOrderByExpressions + i] =
+ DataSchema.ColumnDataType.fromDataType(columnContext.getDataType(), columnContext.isSingleValue());
+ }
+ return new DataSchema(columnNames, columnDataTypes);
+ }
+
+ @Override
+ public String toExplainString() {
+ StringBuilder stringBuilder = new StringBuilder(EXPLAIN_NAME).append("(selectList:");
+ if (!_expressions.isEmpty()) {
+ stringBuilder.append(_expressions.get(0));
+ for (int i = 1; i < _expressions.size(); i++) {
+ stringBuilder.append(", ").append(_expressions.get(i));
+ }
+ }
+ return stringBuilder.append(')').toString();
+ }
+
+ @Override
+ protected String getExplainName() {
+ return CaseFormat.UPPER_UNDERSCORE.to(CaseFormat.UPPER_CAMEL, EXPLAIN_NAME);
+ }
+
+ @Override
+ protected void explainAttributes(ExplainAttributeBuilder attributeBuilder) {
+ super.explainAttributes(attributeBuilder);
+ if (_expressions.isEmpty()) {
+ return;
+ }
+ attributeBuilder.putStringList("selectList",
+ _expressions.stream().map(ExpressionContext::toString).collect(Collectors.toList()));
+ }
+
+ @Override
+ public List<Operator> getChildOperators() {
+ return Collections.singletonList(_projectOperator);
+ }
+
+ @Override
+ public IndexSegment getIndexSegment() {
+ return _indexSegment;
+ }
+
+ @Override
+ public ExecutionStatistics getExecutionStatistics() {
+ long numEntriesScannedInFilter = _projectOperator.getExecutionStatistics().getNumEntriesScannedInFilter();
+ int numTotalDocs = _indexSegment.getSegmentMetadata().getTotalDocs();
+ return new ExecutionStatistics(_numDocsScanned, numEntriesScannedInFilter, _numEntriesScannedPostFilter,
+ numTotalDocs);
+ }
+
+ /// Reports scan cost to the shared {@link QueryScanCostContext} so scan-based query killing
+ /// ({@link org.apache.pinot.core.common.Operator#nextBlock()} -> {@code checkScanBasedKilling}) can see this
+ /// operator's work. Without this the killer is blind on the streaming path, unlike every other selection operator.
+ private void reportScanCost(int numDocsScanned, long numEntriesScannedPostFilter) {
+ QueryScanCostContext scanCost = getScanCostContext();
+ if (scanCost != null) {
+ if (numDocsScanned > 0) {
+ scanCost.addDocsScanned(numDocsScanned);
+ }
+ if (numEntriesScannedPostFilter > 0) {
+ scanCost.addEntriesScannedPostFilter(numEntriesScannedPostFilter);
+ }
+ }
+ }
+}
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/plan/CombinePlanNode.java b/pinot-core/src/main/java/org/apache/pinot/core/plan/CombinePlanNode.java
index 71539035f53..298897e3ec2 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/plan/CombinePlanNode.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/plan/CombinePlanNode.java
@@ -34,6 +34,7 @@ import org.apache.pinot.core.operator.combine.SelectionOnlyCombineOperator;
import org.apache.pinot.core.operator.combine.SelectionOrderByCombineOperator;
import org.apache.pinot.core.operator.combine.SequentialSortedGroupByCombineOperator;
import org.apache.pinot.core.operator.combine.SortedGroupByCombineOperator;
+import org.apache.pinot.core.operator.combine.StreamingSelectionOrderByCombineOperator;
import org.apache.pinot.core.operator.streaming.StreamingDistinctCombineOperator;
import org.apache.pinot.core.operator.streaming.StreamingGroupByCombineOperator;
import org.apache.pinot.core.operator.streaming.StreamingSelectionOnlyCombineOperator;
@@ -131,6 +132,17 @@ public class CombinePlanNode implements PlanNode {
// Use streaming operator only for non-empty selection-only query
return createStreamingSelectionOnlyCombineOperator(operators);
}
+ // 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)
+ && _queryContext.getLimit() != 0) {
+ List<OrderByExpressionContext> orderByExpressions = _queryContext.getOrderByExpressions();
+ if (orderByExpressions != null
+ && orderByExpressions.get(0).getExpression().getType() == ExpressionContext.Type.IDENTIFIER) {
+ return new StreamingSelectionOrderByCombineOperator(operators, _queryContext, _executorService,
+ true /* streaming */);
+ }
+ }
// Streaming flushes partial results, so it needs an aggregation above to merge them back together.
// Leaves that must return final results are excluded, see StreamingGroupByCombineOperator.
boolean leafReturnsFinalResult =
@@ -186,6 +198,10 @@ public class CombinePlanNode implements PlanNode {
List<OrderByExpressionContext> orderByExpressions = _queryContext.getOrderByExpressions();
assert orderByExpressions != null;
if (orderByExpressions.get(0).getExpression().getType() == ExpressionContext.Type.IDENTIFIER) {
+ if (_queryContext.isSortedSelectionMergeEnabled()) {
+ return new StreamingSelectionOrderByCombineOperator(operators, _queryContext, _executorService,
+ false /* streaming */);
+ }
return new MinMaxValueBasedSelectionOrderByCombineOperator(operators, _queryContext, _executorService);
}
return new SelectionOrderByCombineOperator(operators, _queryContext, _executorService);
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/plan/SelectionPlanNode.java b/pinot-core/src/main/java/org/apache/pinot/core/plan/SelectionPlanNode.java
index 029f4811c44..4849a193eeb 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/plan/SelectionPlanNode.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/plan/SelectionPlanNode.java
@@ -32,6 +32,7 @@ import org.apache.pinot.core.operator.query.SelectionOnlyOperator;
import org.apache.pinot.core.operator.query.SelectionOrderByOperator;
import org.apache.pinot.core.operator.query.SelectionPartiallyOrderedByDescOperation;
import org.apache.pinot.core.operator.query.SelectionPartiallyOrderedByLinearOperator;
+import org.apache.pinot.core.operator.query.StreamingSelectionOrderByOperator;
import org.apache.pinot.core.query.request.context.QueryContext;
import org.apache.pinot.core.query.selection.SelectionOperatorUtils;
import org.apache.pinot.segment.spi.IndexSegment;
@@ -88,11 +89,36 @@ public class SelectionPlanNode implements PlanNode {
maxDocsPerCall = Math.min(limit + _queryContext.getOffset(), DocIdSetPlanNode.MAX_DOC_PER_CALL);
}
- BaseProjectOperator<?> projectOperator = getSortedByProject(expressions, maxDocsPerCall, orderByExpressions);
boolean asc = orderByExpressions.get(0).isAsc();
// Remember that we cannot use asc == projectOperator.isAscending() because empty operators are considered
// both ascending and descending
DocIdOrderedOperator.DocIdOrder queryOrder = DocIdOrderedOperator.DocIdOrder.fromAsc(asc);
+
+ // Opt-in streaming path: emit one globally-sorted block at a time so a downstream k-way-merge combine can pull
+ // lazily. Only build it when the first order-by column is an identifier (kept consistent with the combine-side
+ // gate) and the forward-scan project is order-compatible; the DESC-incompatible sorted case still falls back to
+ // the materialized SelectionPartiallyOrderedByDescOperation below so global order stays correct.
+ if (_queryContext.isSortedSelectionMergeEnabled()
+ && orderByExpressions.get(0).getExpression().getType() == ExpressionContext.Type.IDENTIFIER) {
+ // When there are non-order-by output expressions, only fetch the order-by expressions during the forward scan
+ // (the streaming operator fetches the rest in a second pass); otherwise fetch all expressions.
+ List<ExpressionContext> projectExpressions = expressions;
+ if (expressions.size() > numOrderByExpressions) {
+ projectExpressions = new ArrayList<>(numOrderByExpressions);
+ for (OrderByExpressionContext orderByExpression : orderByExpressions) {
+ projectExpressions.add(orderByExpression.getExpression());
+ }
+ }
+ BaseProjectOperator<?> streamingProjectOperator =
+ getSortedByProject(projectExpressions, maxDocsPerCall, orderByExpressions);
+ if (streamingProjectOperator.isCompatibleWith(queryOrder)) {
+ return new StreamingSelectionOrderByOperator(_indexSegment, _queryContext, expressions,
+ streamingProjectOperator, sortedColumnsPrefixSize);
+ }
+ // DESC-incompatible: fall through to the materialized fallback (rebuilds the project over all expressions).
+ }
+
+ BaseProjectOperator<?> projectOperator = getSortedByProject(expressions, maxDocsPerCall, orderByExpressions);
if (projectOperator.isCompatibleWith(queryOrder)) {
return new SelectionPartiallyOrderedByLinearOperator(_indexSegment, _queryContext, expressions, projectOperator,
sortedColumnsPrefixSize);
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java b/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java
index 3a24e0ed51e..25ff2a4db17 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java
@@ -270,6 +270,17 @@ public class InstancePlanMakerImplV2 implements PlanMaker {
}
queryContext.setMaxExecutionThreads(maxExecutionThreads);
+ // Set streaming selection order-by options (opt-in; gated on selection queries to prevent accidental routing
+ // if a downstream guard is ever missed)
+ if (QueryContextUtils.isSelectionQuery(queryContext)) {
+ queryContext.setSortedSelectionMergeEnabled(QueryOptionsUtils.isSortedSelectionMergeEnabled(queryOptions));
+ Integer sortedSelectionMergeBlockSize =
+ QueryOptionsUtils.getSortedSelectionMergeBlockSize(queryOptions);
+ if (sortedSelectionMergeBlockSize != null) {
+ queryContext.setSortedSelectionMergeBlockSize(sortedSelectionMergeBlockSize);
+ }
+ }
+
// Set group-by query options
if (QueryContextUtils.isAggregationQuery(queryContext) && queryContext.getGroupByExpressions() != null) {
// Set maxInitialResultHolderCapacity
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java b/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java
index 1371fd3f3c7..0d7f17775ce 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java
@@ -44,6 +44,7 @@ import org.apache.pinot.core.util.MemoizedClassAssociation;
import org.apache.pinot.segment.spi.datasource.DataSource;
import org.apache.pinot.spi.config.table.FieldConfig;
import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.utils.CommonConstants.Broker;
import org.apache.pinot.spi.utils.CommonConstants.Server;
@@ -135,6 +136,11 @@ public class QueryContext {
private int _streamingGroupByFlushThreshold;
// Flush threshold for streaming distinct (0 = disabled)
private int _streamingDistinctFlushThreshold;
+ /// Opt-in: use the streaming k-way-merge selection ORDER BY combine over sorted segments
+ private boolean _sortedSelectionMergeEnabled;
+ /// Output block size (rows) for the streaming selection ORDER BY combine
+ private int _sortedSelectionMergeBlockSize = Broker.DEFAULT_SORTED_SELECTION_MERGE_BLOCK_SIZE;
+
// Whether null handling is enabled
private boolean _nullHandlingEnabled;
// Whether server returns the final result
@@ -509,6 +515,22 @@ public class QueryContext {
_streamingDistinctFlushThreshold = streamingDistinctFlushThreshold;
}
+ public boolean isSortedSelectionMergeEnabled() {
+ return _sortedSelectionMergeEnabled;
+ }
+
+ public void setSortedSelectionMergeEnabled(boolean sortedSelectionMergeEnabled) {
+ _sortedSelectionMergeEnabled = sortedSelectionMergeEnabled;
+ }
+
+ public int getSortedSelectionMergeBlockSize() {
+ return _sortedSelectionMergeBlockSize;
+ }
+
+ public void setSortedSelectionMergeBlockSize(int sortedSelectionMergeBlockSize) {
+ _sortedSelectionMergeBlockSize = sortedSelectionMergeBlockSize;
+ }
+
public boolean isNullHandlingEnabled() {
return _nullHandlingEnabled;
}
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/CombineSlowOperatorsTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/CombineSlowOperatorsTest.java
index 32a2e9c348b..115ea33a0e5 100644
--- a/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/CombineSlowOperatorsTest.java
+++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/CombineSlowOperatorsTest.java
@@ -24,6 +24,7 @@ import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Supplier;
@@ -41,6 +42,7 @@ import org.apache.pinot.segment.spi.datasource.DataSourceMetadata;
import org.apache.pinot.spi.exception.EarlyTerminationException;
import org.apache.pinot.spi.exception.QueryErrorCode;
import org.apache.pinot.spi.exception.QueryErrorMessage;
+import org.apache.pinot.spi.query.QueryThreadContext;
import org.apache.pinot.util.TestUtils;
import org.testng.Assert;
import org.testng.annotations.AfterClass;
@@ -54,6 +56,7 @@ import static org.mockito.Mockito.when;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertFalse;
import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertTrue;
/// This test mimic the behavior of combining slow operators, where operation is not done by the timeout. When the
@@ -131,16 +134,7 @@ public class CombineSlowOperatorsTest {
@Test
public void testCancelMinMaxValueBasedSelectionOrderByCombineOperator() {
CountDownLatch ready = new CountDownLatch(1);
- List<Operator> operators = getOperators(ready, () -> {
- IndexSegment seg = mock(IndexSegment.class);
- DataSource ds = mock(DataSource.class);
- DataSourceMetadata dsmd = mock(DataSourceMetadata.class);
- when(dsmd.getMinValue()).thenReturn(100L);
- when(dsmd.getMaxValue()).thenReturn(200L);
- when(seg.getDataSource(anyString(), any())).thenReturn(ds);
- when(ds.getDataSourceMetadata()).thenReturn(dsmd);
- return seg;
- });
+ List<Operator> operators = getOperators(ready, minMaxSegmentSupplier());
QueryContext queryContext = QueryContextConverterUtils.getQueryContext("SELECT * FROM testTable ORDER BY column");
queryContext.setEndTimeMs(System.currentTimeMillis() + 10000);
MinMaxValueBasedSelectionOrderByCombineOperator combineOperator =
@@ -170,13 +164,85 @@ public class CombineSlowOperatorsTest {
testCancelCombineOperator(combineOperator, ready);
}
+ @Test
+ public void testCancelStreamingSelectionOrderByCombineOperator() {
+ CountDownLatch ready = new CountDownLatch(1);
+ List<Operator> operators = getOperators(ready, minMaxSegmentSupplier());
+ QueryContext queryContext = QueryContextConverterUtils.getQueryContext("SELECT * FROM testTable ORDER BY column");
+ queryContext.setEndTimeMs(System.currentTimeMillis() + 10000);
+ // Single-stage mode: nextBlock() drives the whole merge synchronously on the (cancellable) caller thread.
+ StreamingSelectionOrderByCombineOperator combineOperator =
+ new StreamingSelectionOrderByCombineOperator(operators, queryContext, _executorService, false);
+ testCancelCombineOperator(combineOperator, ready, operators);
+ }
+
+ @Test
+ public void testCancelStreamingSelectionOrderByCombineOperatorStreamingMode() {
+ CountDownLatch ready = new CountDownLatch(1);
+ List<Operator> operators = getOperators(ready, minMaxSegmentSupplier());
+ QueryContext queryContext = QueryContextConverterUtils.getQueryContext("SELECT * FROM testTable ORDER BY column");
+ queryContext.setEndTimeMs(System.currentTimeMillis() + 10000);
+ // Streaming (MSE-leaf) mode: the first getNextBlock() still drives the merge on the caller thread, so the same
+ // interrupt-on-cancel path applies and must surface an ExceptionResultsBlock.
+ StreamingSelectionOrderByCombineOperator combineOperator =
+ new StreamingSelectionOrderByCombineOperator(operators, queryContext, _executorService, true);
+ testCancelCombineOperator(combineOperator, ready, operators);
+ }
+
+ /// The merge loop drains its heap on the caller thread and, for single-block cursors, may never re-enter a child
+ /// operator, so it must check the query deadline itself. With an already-expired deadline the operator must surface a
+ /// timeout <b>before</b> activating any child - asserted via {@code _operationInProgress}, which also keeps the test
+ /// from passing vacuously if the timeout came from somewhere else.
+ @Test
+ public void testStreamingSelectionOrderByCombineOperatorHonorsDeadline() {
+ List<Operator> operators = getOperators(null, minMaxSegmentSupplier());
+ QueryContext queryContext = QueryContextConverterUtils.getQueryContext("SELECT * FROM testTable ORDER BY column");
+ queryContext.setEndTimeMs(System.currentTimeMillis() - 1);
+ StreamingSelectionOrderByCombineOperator combineOperator =
+ new StreamingSelectionOrderByCombineOperator(operators, queryContext, _executorService, false);
+ try (QueryThreadContext ignore = QueryThreadContext.openForSseTest()) {
+ BaseResultsBlock resultsBlock = combineOperator.nextBlock();
+ assertTrue(resultsBlock instanceof ExceptionResultsBlock,
+ "Expired deadline must surface as an ExceptionResultsBlock, got: " + resultsBlock.getClass().getName());
+ }
+ for (Operator operator : operators) {
+ assertFalse(((SlowOperator) operator)._operationInProgress.get(),
+ "Deadline must be observed before any child operator is driven");
+ }
+ }
+
+ /// A segment whose first order-by column reports a non-null min/max, as required by the min/max-based operators.
+ private static Supplier<IndexSegment> minMaxSegmentSupplier() {
+ return () -> {
+ IndexSegment seg = mock(IndexSegment.class);
+ DataSource ds = mock(DataSource.class);
+ DataSourceMetadata dsmd = mock(DataSourceMetadata.class);
+ when(dsmd.getMinValue()).thenReturn(100L);
+ when(dsmd.getMaxValue()).thenReturn(200L);
+ when(seg.getDataSource(anyString(), any())).thenReturn(ds);
+ when(ds.getDataSourceMetadata()).thenReturn(dsmd);
+ return seg;
+ };
+ }
+
private void testCancelCombineOperator(BaseCombineOperator<?> combineOperator, CountDownLatch ready) {
+ testCancelCombineOperator(combineOperator, ready, List.of());
+ }
+
+ /// Submits the combine operator on a separate thread, waits for a child operator to start, cancels the future, and
+ /// asserts the operator surfaces an {@link ExceptionResultsBlock}. When {@code operatorsToVerify} is non-empty
+ /// (the single-threaded streaming combine, whose merge runs the children on the caller thread), it additionally
+ /// asserts the cancellation was genuine - a child actually started and none completed normally - so the test
+ /// cannot pass vacuously on an unrelated failure or by never exercising the cancel path.
+ private void testCancelCombineOperator(BaseCombineOperator<?> combineOperator, CountDownLatch ready,
+ List<Operator> operatorsToVerify) {
AtomicReference<BaseResultsBlock> resultsBlock = new AtomicReference<>();
// Avoid early finalization by not using Executors.newSingleThreadExecutor (java <= 20, JDK-8145304)
ExecutorService combineExecutor = Executors.newFixedThreadPool(1);
try {
Future<?> future = combineExecutor.submit(() -> resultsBlock.set(combineOperator.nextBlock()));
- ready.await();
+ // Bound the wait so a regression where no child operator ever starts fails fast here instead of hanging.
+ assertTrue(ready.await(10, TimeUnit.SECONDS), "Expected a child operator to start before cancellation");
// At this point, the combineOperator is or will be waiting on future.get() for all sub operators, and the
// waiting can be cancelled as below.
future.cancel(true);
@@ -187,6 +253,12 @@ public class CombineSlowOperatorsTest {
}
TestUtils.waitForCondition((aVoid) -> resultsBlock.get() instanceof ExceptionResultsBlock, 10_000,
"Should have been cancelled");
+ // Genuine-cancellation check: no child ran to normal completion (each was interrupted or never started), so the
+ // ExceptionResultsBlock above is the cancellation outcome rather than an unrelated error.
+ for (Operator operator : operatorsToVerify) {
+ assertFalse(((SlowOperator) operator)._notInterrupted.get(),
+ "No operator should have completed normally after cancellation");
+ }
}
/// NOTE: It is hard to test the logger behavior, but only one error message about the query timeout should be logged
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java
new file mode 100644
index 00000000000..e2d9c9f8a9b
--- /dev/null
+++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java
@@ -0,0 +1,553 @@
+/**
+ * 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.io.File;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.List;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.stream.Collectors;
+import org.apache.commons.io.FileUtils;
+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.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.plan.CombinePlanNode;
+import org.apache.pinot.core.plan.PlanNode;
+import org.apache.pinot.core.plan.maker.InstancePlanMakerImplV2;
+import org.apache.pinot.core.plan.maker.PlanMaker;
+import org.apache.pinot.core.query.executor.ResultsBlockStreamer;
+import org.apache.pinot.core.query.request.context.QueryContext;
+import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils;
+import org.apache.pinot.core.query.utils.OrderByComparatorFactory;
+import org.apache.pinot.core.util.QueryMultiThreadingUtils;
+import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader;
+import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl;
+import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader;
+import org.apache.pinot.segment.spi.IndexSegment;
+import org.apache.pinot.segment.spi.SegmentContext;
+import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.data.readers.GenericRow;
+import org.apache.pinot.spi.utils.CommonConstants.Server;
+import org.apache.pinot.spi.utils.ReadMode;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.intellij.lang.annotations.Language;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertTrue;
+
+
+/// Combine-level tests for {@link StreamingSelectionOrderByCombineOperator} (step-3 operator) and its wiring into
+/// {@link CombinePlanNode#getCombineOperator()} (step-4).
+///
+/// <p>The streaming combine must return the same globally-sorted top-K rows as the default
+/// {@link MinMaxValueBasedSelectionOrderByCombineOperator}, only (in streaming mode) spread across several bounded
+/// blocks. Each functional test therefore asserts <b>streaming-vs-non-streaming parity</b>: it runs the identical query
+/// twice over the same in-memory segments - once with {@code sortedSelectionMergeEnabled=true} (asserting the new
+/// operator was actually selected) and once with the hint off (asserting the {@code MinMax} operator was selected) -
+/// then checks the two row sets are equal as a multiset and that the streaming output is fully sorted by the order-by
+/// comparator.
+///
+/// <p>To keep the top-K boundary unambiguous (operators may legitimately disagree on which of several rows that tie on
+/// every order-by key fall inside the limit) every parity query ends its ORDER BY with the globally-unique
+/// {@code valCol}
+/// so the comparator is a total order; the merge still genuinely interleaves segments because the primary sort column
+/// ({@code sortedCol}) overlaps across segments. Multiset (rather than positional) comparison then tolerates only the
+/// harmless reordering of fully-equal projected rows.
+public class StreamingSelectionOrderByCombineOperatorTest {
+ private static final File TEMP_DIR =
+ new File(FileUtils.getTempDirectory(), "StreamingSelectionOrderByCombineOperatorTest");
+ private static final String RAW_TABLE_NAME = "testTable";
+
+ private static final String SORTED_COL = "sortedCol";
+ private static final String TAIL_COL = "tailCol";
+ private static final String VAL_COL = "valCol";
+ private static final String NULLABLE_COL = "nullableCol";
+ /// Non-INT projected columns so the parity assertion can catch a stored-type / boxing regression (e.g. LONG emitted
+ /// where MinMax emits INT), which an all-INT suite cannot observe.
+ private static final String LONG_COL = "longCol";
+ private static final String STR_COL = "strCol";
+
+ /// Create (MAX_NUM_THREADS_PER_QUERY * 2) sorted segments so the leaf runs plan nodes across multiple threads.
+ private static final int NUM_SEGMENTS = QueryMultiThreadingUtils.MAX_NUM_THREADS_PER_QUERY * 2;
+ private static final int NUM_RECORDS_PER_SEGMENT = 100;
+
+ private static final TableConfig SORTED_TABLE_CONFIG =
+ new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME).setSortedColumn(SORTED_COL).build();
+ private static final TableConfig UNSORTED_TABLE_CONFIG =
+ new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME).build();
+ private static final Schema SCHEMA = new Schema.SchemaBuilder()
+ .addSingleValueDimension(SORTED_COL, FieldSpec.DataType.INT)
+ .addSingleValueDimension(TAIL_COL, FieldSpec.DataType.INT)
+ .addSingleValueDimension(VAL_COL, FieldSpec.DataType.INT)
+ .addSingleValueDimension(NULLABLE_COL, FieldSpec.DataType.INT)
+ .addSingleValueDimension(LONG_COL, FieldSpec.DataType.LONG)
+ .addSingleValueDimension(STR_COL, FieldSpec.DataType.STRING)
+ .build();
+
+ private static final PlanMaker PLAN_MAKER = new InstancePlanMakerImplV2();
+ private static final ExecutorService EXECUTOR = Executors.newCachedThreadPool();
+
+ /// Sorted segments with overlapping primary-column ranges, so the k-way merge interleaves them (genuine merge rather
+ /// than concatenation). Built with null handling on so the null-handling test sees real nulls in NULLABLE_COL; reads
+ /// with null handling off fall back to the column default.
+ private List<IndexSegment> _sortedSegments;
+ /// Sorted segments with disjoint, globally-increasing ranges: ORDER BY sortedCol with a small LIMIT drains only the
+ /// lowest segment, so min/max pruning must skip the rest (none acquired/scanned).
+ private List<IndexSegment> _disjointSegments;
+ /// A mix of sorted (streaming child) and physically-unsorted (single materialized top-K block child) segments,
+ /// exercising both SegmentCursor backings in one merge.
+ private List<IndexSegment> _mixedSegments;
+ /// Very low cardinality primary column (4 distinct values across 100 rows) so each value is a long run: exercises the
+ /// run/heap path in the streaming children and ties on the primary key at the prune boundary.
+ private List<IndexSegment> _lowCardSegments;
+
+ @BeforeClass
+ public void setUp()
+ throws Exception {
+ FileUtils.deleteDirectory(TEMP_DIR);
+
+ _sortedSegments = new ArrayList<>(NUM_SEGMENTS);
+ for (int i = 0; i < NUM_SEGMENTS; i++) {
+ _sortedSegments.add(buildSegment(SORTED_TABLE_CONFIG, "sorted_" + i, buildOverlappingSortedRecords(i), true));
+ }
+
+ _disjointSegments = new ArrayList<>(NUM_SEGMENTS);
+ for (int i = 0; i < NUM_SEGMENTS; i++) {
+ _disjointSegments.add(buildSegment(SORTED_TABLE_CONFIG, "disjoint_" + i, buildDisjointSortedRecords(i), false));
+ }
+
+ _lowCardSegments = new ArrayList<>(NUM_SEGMENTS);
+ for (int i = 0; i < NUM_SEGMENTS; i++) {
+ _lowCardSegments.add(buildSegment(SORTED_TABLE_CONFIG, "lowCard_" + i, buildLowCardinalityRecords(i), false));
+ }
+
+ // Two sorted + two unsorted segments, globally-unique valCol across all four so the multiset comparison is exact.
+ _mixedSegments = new ArrayList<>(4);
+ _mixedSegments.add(buildSegment(SORTED_TABLE_CONFIG, "mixedSorted_0", buildOverlappingSortedRecords(0), false));
+ _mixedSegments.add(buildSegment(SORTED_TABLE_CONFIG, "mixedSorted_1", buildOverlappingSortedRecords(1), false));
+ _mixedSegments.add(buildSegment(UNSORTED_TABLE_CONFIG, "mixedUnsorted_0", buildUnsortedRecords(2), false));
+ _mixedSegments.add(buildSegment(UNSORTED_TABLE_CONFIG, "mixedUnsorted_1", buildUnsortedRecords(3), false));
+ }
+
+ private static List<GenericRow> buildOverlappingSortedRecords(int index) {
+ int baseValue = index * NUM_RECORDS_PER_SEGMENT / 2;
+ List<GenericRow> records = new ArrayList<>(NUM_RECORDS_PER_SEGMENT);
+ for (int i = 0; i < NUM_RECORDS_PER_SEGMENT; i++) {
+ GenericRow record = new GenericRow();
+ record.putValue(SORTED_COL, baseValue + i);
+ record.putValue(TAIL_COL, NUM_RECORDS_PER_SEGMENT - i);
+ // Globally unique across all segments -> a total order when used as the final order-by key.
+ record.putValue(VAL_COL, index * 1_000_000 + i);
+ // Beyond the int range so a regression that narrows LONG -> INT would change the boxed value.
+ record.putValue(LONG_COL, 10_000_000_000L + index * 1_000_000L + i);
+ record.putValue(STR_COL, "s_" + index + "_" + i);
+ // Every 7th row is null so the null-handling test exercises null projection.
+ if (i % 7 == 0) {
+ record.addNullValueField(NULLABLE_COL);
+ } else {
+ record.putValue(NULLABLE_COL, i);
+ }
+ records.add(record);
+ }
+ return records;
+ }
+
+ private static List<GenericRow> buildDisjointSortedRecords(int index) {
+ List<GenericRow> records = new ArrayList<>(NUM_RECORDS_PER_SEGMENT);
+ int baseValue = index * 1000;
+ for (int i = 0; i < NUM_RECORDS_PER_SEGMENT; i++) {
+ GenericRow record = new GenericRow();
+ record.putValue(SORTED_COL, baseValue + i);
+ record.putValue(TAIL_COL, i);
+ record.putValue(VAL_COL, baseValue + i);
+ record.putValue(NULLABLE_COL, i);
+ record.putValue(LONG_COL, 10_000_000_000L + baseValue + i);
+ record.putValue(STR_COL, "d_" + index + "_" + i);
+ records.add(record);
+ }
+ return records;
+ }
+
+ private static List<GenericRow> buildLowCardinalityRecords(int index) {
+ List<GenericRow> records = new ArrayList<>(NUM_RECORDS_PER_SEGMENT);
+ for (int i = 0; i < NUM_RECORDS_PER_SEGMENT; i++) {
+ GenericRow record = new GenericRow();
+ // 4 distinct values per segment, non-decreasing so the segment is physically sorted on SORTED_COL.
+ record.putValue(SORTED_COL, i / 25);
+ record.putValue(TAIL_COL, NUM_RECORDS_PER_SEGMENT - i);
+ record.putValue(VAL_COL, index * 1_000_000 + i);
+ record.putValue(NULLABLE_COL, i);
+ record.putValue(LONG_COL, 10_000_000_000L + index * 1_000_000L + i);
+ record.putValue(STR_COL, "l_" + index + "_" + i);
+ records.add(record);
+ }
+ return records;
+ }
+
+ private static List<GenericRow> buildUnsortedRecords(int index) {
+ List<GenericRow> records = new ArrayList<>(NUM_RECORDS_PER_SEGMENT);
+ for (int i = 0; i < NUM_RECORDS_PER_SEGMENT; i++) {
+ GenericRow record = new GenericRow();
+ // A non-monotonic permutation of [0, NUM_RECORDS_PER_SEGMENT) (7919 is prime and coprime with 100), so the
+ // column is genuinely not physically sorted and SelectionPlanNode falls back to a materialized top-K block.
+ record.putValue(SORTED_COL, (i * 7919) % NUM_RECORDS_PER_SEGMENT);
+ record.putValue(TAIL_COL, i);
+ record.putValue(VAL_COL, index * 1_000_000 + i);
+ record.putValue(NULLABLE_COL, i);
+ record.putValue(LONG_COL, 10_000_000_000L + index * 1_000_000L + i);
+ record.putValue(STR_COL, "u_" + index + "_" + i);
+ records.add(record);
+ }
+ return records;
+ }
+
+ private static IndexSegment buildSegment(TableConfig tableConfig, String segmentName, List<GenericRow> records,
+ boolean nullHandling)
+ throws Exception {
+ SegmentGeneratorConfig segmentGeneratorConfig = new SegmentGeneratorConfig(tableConfig, SCHEMA);
+ segmentGeneratorConfig.setTableName(RAW_TABLE_NAME);
+ segmentGeneratorConfig.setSegmentName(segmentName);
+ segmentGeneratorConfig.setDefaultNullHandlingEnabled(nullHandling);
+ segmentGeneratorConfig.setOutDir(TEMP_DIR.getPath());
+
+ SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl();
+ driver.init(segmentGeneratorConfig, new GenericRowRecordReader(records));
+ driver.build();
+
+ return ImmutableSegmentLoader.load(new File(TEMP_DIR, segmentName), ReadMode.mmap);
+ }
+
+ @Test
+ public void testAscendingParity() {
+ assertParity(_sortedSegments, "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50", false);
+ }
+
+ @Test
+ public void testDescendingParity() {
+ // Reverse order must be allowed for the per-segment forward-scan to iterate sortedCol descending; otherwise the
+ // segment falls back to the materialized DESC operator (covered separately by testDescIncompatibleFallbackParity).
+ assertParity(_sortedSegments,
+ "SET allowReverseOrder=true; SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol DESC, valCol DESC "
+ + "LIMIT 50", false);
+ }
+
+ @Test
+ public void testDescIncompatibleFallbackParity() {
+ // allowReverseOrder=false + DESC -> the streaming child cannot scan descending, so SelectionPlanNode emits the
+ // materialized DESC top-K block; the combine still routes to the streaming combine and merges single-block cursors.
+ assertParity(_sortedSegments,
+ "SET allowReverseOrder=false; SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol DESC, valCol DESC "
+ + "LIMIT 50", false);
+ }
+
+ @Test
+ public void testLimitOffsetParity() {
+ // The server retains limit + offset rows; the broker applies the offset later, so both operators keep 40 rows.
+ assertParity(_sortedSegments,
+ "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 30 OFFSET 10", false);
+ }
+
+ @Test
+ public void testLowCardinalityMultiColumnParity() {
+ // Low-cardinality primary column => long runs and many sortedCol ties at the prune boundary; valCol breaks ties.
+ assertParity(_lowCardSegments,
+ "SELECT sortedCol, tailCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 40", false);
+ }
+
+ @Test
+ public void testTwoPhaseSelectNonOrderByParity() {
+ // tailCol is selected but not an order-by key -> the streaming children take the two-phase (order-by-then-fetch)
+ // path. valCol is order-by-only, exercising the phase-1 projection of a non-selected order-by column.
+ assertParity(_sortedSegments, "SELECT tailCol, sortedCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50",
+ false);
+ }
+
+ @Test
+ public void testNonIntProjectionParity() {
+ // Projects a LONG and a STRING column so the multiset comparison would catch a stored-type / boxing regression that
+ // an all-INT projection cannot observe.
+ assertParity(_sortedSegments,
+ "SELECT strCol, longCol, sortedCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50", false);
+ }
+
+ @Test
+ public void testMixedSortedAndUnsortedSegmentsParity() {
+ assertParity(_mixedSegments, "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50", false);
+ }
+
+ @Test
+ public void testNullHandlingEnabledParity() {
+ // Null handling on disables min/max pruning (the combine activates every segment); nullableCol carries real nulls.
+ assertParity(_sortedSegments,
+ "SELECT nullableCol, sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50", true);
+ }
+
+ @Test
+ public void testNullHandlingDisabledParity() {
+ // Same segments/query as the enabled case but null handling off: nulls read back as the column default, pruning on.
+ assertParity(_sortedSegments,
+ "SELECT nullableCol, sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50", false);
+ }
+
+ @Test
+ public void testStreamingMultiBlockExactCount() {
+ // A precise check on the bounded-flush behavior: 20 rows flushed in blocks of 3 yields ceil(20/3) = 7 data blocks.
+ int blockSize = 3;
+ int limit = 20;
+ @Language("sql") String query =
+ "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT " + limit;
+ Result streaming = run(_sortedSegments, query, true, false, true, blockSize);
+ assertEquals(streaming._combineOperator.getClass(), StreamingSelectionOrderByCombineOperator.class);
+ assertEquals(streaming._rows.size(), limit);
+ assertEquals(streaming._numBlocks, (limit + blockSize - 1) / blockSize, "Unexpected number of streamed blocks");
+ assertSorted(streaming._rows, orderByComparator(query, false));
+ assertMultisetEquals(streaming._rows, run(_sortedSegments, query, false, false, false, 0)._rows);
+ }
+
+ @Test
+ public void testPruningSkipsOutOfTopKSegments() {
+ // Disjoint, globally-increasing ranges + small LIMIT: only the lowest segment can contribute, the rest are pruned
+ // (never acquired or scanned), so far fewer than all docs are scanned.
+ Result result = run(_disjointSegments, "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 5",
+ true, false, false, 0);
+ assertTrue(result._combineOperator instanceof StreamingSelectionOrderByCombineOperator);
+ assertEquals(result._rows.size(), 5);
+ for (int i = 0; i < 5; i++) {
+ assertEquals((int) result._rows.get(i)[0], i, "Unexpected value at position " + i);
+ }
+ int totalDocs = NUM_SEGMENTS * NUM_RECORDS_PER_SEGMENT;
+ assertTrue(result._numDocsScanned < totalDocs, "Pruning should avoid scanning every doc, scanned: "
+ + result._numDocsScanned + " of " + totalDocs);
+ assertTrue(result._numDocsScanned <= NUM_RECORDS_PER_SEGMENT,
+ "Only the lowest-range segment should be scanned, but docs scanned was: " + result._numDocsScanned);
+ }
+
+ @Test
+ public void testEmptyResultSchemaFallback() {
+ // A filter that matches nothing: every streaming child returns no rows, so the combine rebuilds the result schema
+ // from the segment metadata (order-by expressions first) rather than from a child block.
+ Result result = run(_sortedSegments,
+ "SELECT sortedCol, valCol FROM testTable WHERE sortedCol < 0 ORDER BY sortedCol, valCol LIMIT 10", true, false,
+ false, 0);
+ assertTrue(result._combineOperator instanceof StreamingSelectionOrderByCombineOperator);
+ assertTrue(result._rows.isEmpty(), "Expected an empty result, got: " + result._rows.size() + " rows");
+ assertEquals(result._schema, new DataSchema(new String[]{SORTED_COL, VAL_COL},
+ new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.INT}));
+ }
+
+ @Test
+ public void testHintOffSelectsMinMaxOperator() {
+ // Default behavior is unchanged when the hint is off: the classic MinMax operator is still selected.
+ Result result = run(_sortedSegments, "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50",
+ false, false, false, 0);
+ assertTrue(result._combineOperator instanceof MinMaxValueBasedSelectionOrderByCombineOperator,
+ "Hint off must keep the default MinMax combine operator, got: "
+ + result._combineOperator.getClass().getSimpleName());
+ }
+
+ @Test
+ public void testNonIdentifierOrderByFallsBack() {
+ // Even with the hint on, a non-identifier first order-by expression falls back to SelectionOrderByCombineOperator
+ // (the streaming operator and its segment-level counterpart only support a leading identifier).
+ Result result = run(_sortedSegments,
+ "SELECT sortedCol, valCol FROM testTable ORDER BY ADD(sortedCol, 1), valCol LIMIT 50", true, false, false, 0);
+ assertEquals(result._combineOperator.getClass(), SelectionOrderByCombineOperator.class,
+ "Non-identifier first order-by must fall back to SelectionOrderByCombineOperator, got: "
+ + result._combineOperator.getClass().getSimpleName());
+ }
+
+ /// Asserts streaming-vs-non-streaming parity for {@code query}. Runs the MinMax combine (hint off) as the reference,
+ /// then runs the streaming combine in BOTH single-block mode and bounded multi-block streaming mode, asserting each
+ /// selects the streaming operator and produces rows that are sorted by the order-by comparator and equal the MinMax
+ /// rows as a multiset. The streaming variant additionally checks the bounded-flush invariants.
+ private void assertParity(List<IndexSegment> segments, @Language("sql") String query, boolean nullHandling) {
+ Result baseline = run(segments, query, false, nullHandling, false, 0);
+ assertEquals(baseline._combineOperator.getClass(), MinMaxValueBasedSelectionOrderByCombineOperator.class,
+ "Baseline must be the MinMax combine operator, got: " + baseline._combineOperator.getClass().getSimpleName());
+ Comparator<Object[]> comparator = orderByComparator(query, nullHandling);
+
+ // Classic single-stage path (null streamer): the streaming combine flushes the whole merge as one block.
+ Result singleStage = run(segments, query, true, nullHandling, false, 0);
+ assertStreamingParity(singleStage, baseline, comparator, query, false, 0);
+
+ // MSE leaf path (non-null streamer): a small block size forces several bounded data blocks before the metadata
+ // block, genuinely exercising the streaming flush path rather than a single trimmed block.
+ int blockSize = 3;
+ Result streamed = run(segments, query, true, nullHandling, true, blockSize);
+ assertStreamingParity(streamed, baseline, comparator, query, true, blockSize);
+ }
+
+ private void assertStreamingParity(Result result, Result baseline, Comparator<Object[]> comparator,
+ @Language("sql") String query, boolean streaming, int blockSize) {
+ assertEquals(result._combineOperator.getClass(), StreamingSelectionOrderByCombineOperator.class,
+ "Expected the streaming combine operator for query: " + query);
+ assertEquals(result._schema, baseline._schema, "Schema mismatch for query: " + query);
+ assertSorted(result._rows, comparator);
+ assertMultisetEquals(result._rows, baseline._rows);
+ if (streaming) {
+ int total = 0;
+ for (int size : result._blockSizes) {
+ assertTrue(size > 0 && size <= blockSize,
+ "Streamed block size out of range (0, " + blockSize + "] for query " + query + ": " + size);
+ total += size;
+ }
+ assertEquals(total, result._rows.size(), "Streamed block sizes must sum to the row count for query: " + query);
+ if (result._rows.size() > blockSize) {
+ assertTrue(result._numBlocks >= 2, "Expected multiple streamed blocks for query: " + query);
+ }
+ }
+ }
+
+ /// Runs one combine over {@code segments} and collects its rows, blocks, schema and docs-scanned stat.
+ private Result run(List<IndexSegment> segments, @Language("sql") String query, boolean hintOn, boolean nullHandling,
+ boolean streaming, int blockSize) {
+ QueryContext queryContext = QueryContextConverterUtils.getQueryContext(query);
+ queryContext.setNullHandlingEnabled(nullHandling);
+ if (hintOn) {
+ queryContext.setSortedSelectionMergeEnabled(true);
+ if (blockSize > 0) {
+ queryContext.setSortedSelectionMergeBlockSize(blockSize);
+ }
+ }
+ queryContext.setEndTimeMs(System.currentTimeMillis() + Server.DEFAULT_QUERY_EXECUTOR_TIMEOUT_MS);
+
+ List<PlanNode> planNodes = new ArrayList<>(segments.size());
+ for (IndexSegment segment : segments) {
+ SegmentContext segmentContext = new SegmentContext(segment);
+ planNodes.add(streaming ? PLAN_MAKER.makeStreamingSegmentPlanNode(segmentContext, queryContext)
+ : PLAN_MAKER.makeSegmentPlanNode(segmentContext, queryContext));
+ }
+ ResultsBlockStreamer streamer = streaming ? block -> {
+ } : null;
+ CombinePlanNode combinePlanNode = new CombinePlanNode(planNodes, queryContext, EXECUTOR, streamer);
+
+ Result result = new Result();
+ Operator<?> combineOperator = combinePlanNode.run();
+ result._combineOperator = combineOperator;
+ result._rows = new ArrayList<>();
+ result._blockSizes = new ArrayList<>();
+ if (streaming) {
+ // Drive the streaming combine: collect bounded data blocks until the terminal metadata block, which carries the
+ // aggregated execution stats.
+ while (true) {
+ BaseResultsBlock block = (BaseResultsBlock) combineOperator.nextBlock();
+ if (block instanceof MetadataResultsBlock) {
+ if (result._schema == null) {
+ result._schema = block.getDataSchema();
+ }
+ result._numDocsScanned = block.getNumDocsScanned();
+ break;
+ }
+ SelectionResultsBlock dataBlock = (SelectionResultsBlock) block;
+ if (result._schema == null) {
+ result._schema = dataBlock.getDataSchema();
+ }
+ List<Object[]> rows = dataBlock.getRows();
+ assertNotNull(rows);
+ result._rows.addAll(rows);
+ result._blockSizes.add(rows.size());
+ result._numBlocks++;
+ assertTrue(result._numBlocks < 1_000_000, "Streaming combine did not terminate");
+ }
+ } else {
+ SelectionResultsBlock block = (SelectionResultsBlock) combineOperator.nextBlock();
+ result._schema = block.getDataSchema();
+ List<Object[]> rows = block.getRows();
+ assertNotNull(rows);
+ result._rows.addAll(rows);
+ result._blockSizes.add(rows.size());
+ result._numBlocks = 1;
+ result._numDocsScanned = block.getNumDocsScanned();
+ }
+ return result;
+ }
+
+ private static Comparator<Object[]> orderByComparator(@Language("sql") String query, boolean nullHandling) {
+ QueryContext queryContext = QueryContextConverterUtils.getQueryContext(query);
+ List<OrderByExpressionContext> orderByExpressions = queryContext.getOrderByExpressions();
+ assertNotNull(orderByExpressions);
+ return OrderByComparatorFactory.getComparator(orderByExpressions, nullHandling);
+ }
+
+ private static void assertSorted(List<Object[]> rows, Comparator<Object[]> comparator) {
+ for (int i = 1; i < rows.size(); i++) {
+ assertTrue(comparator.compare(rows.get(i - 1), rows.get(i)) <= 0,
+ "Rows not sorted by the order-by comparator at position " + i);
+ }
+ }
+
+ /// Asserts the two row lists contain the same rows, independent of the ordering of fully-equal projected rows.
+ private static void assertMultisetEquals(List<Object[]> actual, List<Object[]> expected) {
+ assertEquals(toCanonical(actual), toCanonical(expected), "Row multisets differ");
+ }
+
+ /// Canonicalizes rows for multiset comparison. Each cell is encoded with its runtime class so a stored-type / boxing
+ /// regression (e.g. a LONG emitted where the reference emits INT) changes the encoding and fails the assertion, which
+ /// a plain {@code Arrays.toString} (type-blind) comparison would miss.
+ private static List<String> toCanonical(List<Object[]> rows) {
+ return rows.stream().map(row -> {
+ StringBuilder sb = new StringBuilder("[");
+ for (int i = 0; i < row.length; i++) {
+ if (i > 0) {
+ sb.append(", ");
+ }
+ Object cell = row[i];
+ sb.append(cell == null ? "null" : cell.getClass().getSimpleName() + ":" + cell);
+ }
+ return sb.append(']').toString();
+ }).sorted().collect(Collectors.toList());
+ }
+
+ @AfterClass
+ public void tearDown()
+ throws IOException {
+ EXECUTOR.shutdownNow();
+ for (List<IndexSegment> segments : List.of(_sortedSegments, _disjointSegments, _mixedSegments, _lowCardSegments)) {
+ for (IndexSegment segment : segments) {
+ segment.destroy();
+ }
+ }
+ FileUtils.deleteDirectory(TEMP_DIR);
+ }
+
+ /// Captured output of a single combine run.
+ private static class Result {
+ private Operator<?> _combineOperator;
+ private DataSchema _schema;
+ private List<Object[]> _rows;
+ private List<Integer> _blockSizes;
+ private int _numBlocks;
+ private long _numDocsScanned;
+ }
+}
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperatorTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperatorTest.java
new file mode 100644
index 00000000000..73a071f673a
--- /dev/null
+++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperatorTest.java
@@ -0,0 +1,406 @@
+/**
+ * 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 java.io.File;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.commons.io.FileUtils;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.core.common.Operator;
+import org.apache.pinot.core.operator.blocks.results.SelectionResultsBlock;
+import org.apache.pinot.core.plan.SelectionPlanNode;
+import org.apache.pinot.core.query.request.context.QueryContext;
+import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils;
+import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader;
+import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl;
+import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader;
+import org.apache.pinot.segment.spi.IndexSegment;
+import org.apache.pinot.segment.spi.SegmentContext;
+import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.data.readers.GenericRow;
+import org.apache.pinot.spi.utils.ReadMode;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.intellij.lang.annotations.Language;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertTrue;
+
+
+/// Segment-level tests for {@link StreamingSelectionOrderByOperator}.
+///
+/// <p>The operator emits the same globally-sorted rows as the existing materialized selection ORDER BY operators, only
+/// spread across many lazily-produced blocks. Each test therefore asserts <b>stream-vs-materialized parity</b>: it
+/// drives the streaming operator through {@link SelectionPlanNode} with {@code sortedSelectionMergeEnabled=true},
+/// concatenates
+/// every {@link Operator#nextBlock()} output until {@code null}, and asserts the concatenation equals the single block
+/// the materialized operator ({@link SelectionPartiallyOrderedByLinearOperator} / {@link SelectionOrderByOperator})
+/// produces for the identical query with the hint off.
+///
+/// <p>To keep element-wise comparison deterministic (priority-queue draining is not stable for rows that tie on every
+/// order-by column) each fixture makes the order-by column tuple unique per row: the {@code _segment} fixture has a
+/// unique sorted column, while the {@code _dupSegment} / {@code _largeSegment} fixtures repeat the sorted column but
+/// pair it with a unique {@code TAIL_COL} order-by tail. The all-order-by-columns-tie case (where stream and
+/// materialized output may legitimately differ in row order) is intentionally out of scope here and is covered by the
+/// combine-level test via multiset comparison.
+public class StreamingSelectionOrderByOperatorTest {
+ private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), "StreamingSelectionOrderByOperatorTest");
+ private static final String RAW_TABLE_NAME = "testTable";
+
+ private static final String SORTED_COL = "sortedCol";
+ private static final String TAIL_COL = "tailCol";
+ private static final String VAL_COL = "valCol";
+ private static final String NULLABLE_COL = "nullableCol";
+
+ private static final TableConfig TABLE_CONFIG =
+ new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME).setSortedColumn(SORTED_COL).build();
+ private static final Schema SCHEMA = new Schema.SchemaBuilder()
+ .addSingleValueDimension(SORTED_COL, FieldSpec.DataType.INT)
+ .addSingleValueDimension(TAIL_COL, FieldSpec.DataType.INT)
+ .addSingleValueDimension(VAL_COL, FieldSpec.DataType.INT)
+ .addSingleValueDimension(NULLABLE_COL, FieldSpec.DataType.INT)
+ .build();
+
+ /// Unique sorted column, no nulls. Exercises the no-tail emission mode.
+ private static final int NUM_RECORDS = 30;
+ /// Repeated sorted column (RUN_SIZE rows per value) paired with a unique tail. Exercises the run / heap path, and
+ /// carries nulls in NULLABLE_COL for the null-handling cases.
+ private static final int NUM_DISTINCT_SORTED = 10;
+ private static final int RUN_SIZE = 4;
+ private static final int NUM_DUP_RECORDS = NUM_DISTINCT_SORTED * RUN_SIZE;
+ /// A single primary-value run larger than one project block (DocIdSetPlanNode.MAX_DOC_PER_CALL == 10_000), so the
+ /// forward scan crosses a block boundary within one run.
+ private static final int NUM_LARGE_RECORDS = 12_000;
+ /// Two runs where the first is larger than one project block: with a limit greater than the first run's size the
+ /// operator both crosses a project-block boundary mid-run AND carries _pendingRow across a run boundary to emit a
+ /// second block.
+ private static final int MULTI_RUN0_SIZE = 10_500;
+ private static final int MULTI_RUN1_SIZE = 1_000;
+ private static final int NUM_MULTI_RUN_RECORDS = MULTI_RUN0_SIZE + MULTI_RUN1_SIZE;
+
+ private IndexSegment _segment;
+ private IndexSegment _dupSegment;
+ private IndexSegment _largeSegment;
+ private IndexSegment _multiRunSegment;
+ /// Carries nulls in an order-by (tail) column so the streaming operator's own null path (phase-1 materialization and
+ /// the null-aware comparator) is exercised, not just the shared two-phase fetch.
+ private IndexSegment _nullTailSegment;
+
+ @BeforeClass
+ public void setUp()
+ throws Exception {
+ FileUtils.deleteDirectory(TEMP_DIR);
+ _segment = buildSegment("uniqueSorted", buildUniqueSortedRecords(), false);
+ _dupSegment = buildSegment("dupSorted", buildDupSortedRecords(), true);
+ _largeSegment = buildSegment("largeRun", buildLargeRunRecords(), false);
+ _multiRunSegment = buildSegment("multiRun", buildMultiRunRecords(), false);
+ _nullTailSegment = buildSegment("nullTail", buildNullTailRecords(), true);
+ }
+
+ private static List<GenericRow> buildUniqueSortedRecords() {
+ List<GenericRow> records = new ArrayList<>(NUM_RECORDS);
+ for (int i = 0; i < NUM_RECORDS; i++) {
+ GenericRow record = new GenericRow();
+ record.putValue(SORTED_COL, i);
+ record.putValue(TAIL_COL, NUM_RECORDS - i);
+ record.putValue(VAL_COL, i * 3);
+ record.putValue(NULLABLE_COL, i);
+ records.add(record);
+ }
+ return records;
+ }
+
+ private static List<GenericRow> buildDupSortedRecords() {
+ List<GenericRow> records = new ArrayList<>(NUM_DUP_RECORDS);
+ for (int i = 0; i < NUM_DUP_RECORDS; i++) {
+ GenericRow record = new GenericRow();
+ record.putValue(SORTED_COL, i / RUN_SIZE);
+ // Tail resets to a descending sequence within each run so the column is NOT globally sorted (the run/heap path is
+ // only taken when the tail is unsorted); the (sortedCol, tailCol) tuple is still unique per row.
+ record.putValue(TAIL_COL, RUN_SIZE - 1 - (i % RUN_SIZE));
+ record.putValue(VAL_COL, i * 2);
+ // Every third row carries a null so the two-phase fetch is exercised with and without null handling.
+ if (i % 3 == 0) {
+ record.addNullValueField(NULLABLE_COL);
+ } else {
+ record.putValue(NULLABLE_COL, i);
+ }
+ records.add(record);
+ }
+ return records;
+ }
+
+ private static List<GenericRow> buildLargeRunRecords() {
+ List<GenericRow> records = new ArrayList<>(NUM_LARGE_RECORDS);
+ for (int i = 0; i < NUM_LARGE_RECORDS; i++) {
+ GenericRow record = new GenericRow();
+ // All rows share one sorted value, so they form a single run spanning multiple project blocks.
+ record.putValue(SORTED_COL, 0);
+ // A non-monotonic permutation of [0, NUM_LARGE_RECORDS) (7919 is prime and coprime with NUM_LARGE_RECORDS, so the
+ // mapping is a bijection): unique tail values that are not physically sorted, forcing the run/heap path.
+ record.putValue(TAIL_COL, (i * 7919) % NUM_LARGE_RECORDS);
+ record.putValue(VAL_COL, i);
+ record.putValue(NULLABLE_COL, i);
+ records.add(record);
+ }
+ return records;
+ }
+
+ private static List<GenericRow> buildMultiRunRecords() {
+ List<GenericRow> records = new ArrayList<>(NUM_MULTI_RUN_RECORDS);
+ appendRun(records, 0, MULTI_RUN0_SIZE);
+ appendRun(records, 1, MULTI_RUN1_SIZE);
+ return records;
+ }
+
+ /// Appends one run of {@code runSize} rows all sharing {@code sortedValue}, with a tail that descends within the run
+ /// (so the tail column is not globally sorted) and is unique per (sortedCol, tailCol) tuple.
+ private static void appendRun(List<GenericRow> records, int sortedValue, int runSize) {
+ for (int j = 0; j < runSize; j++) {
+ GenericRow record = new GenericRow();
+ record.putValue(SORTED_COL, sortedValue);
+ record.putValue(TAIL_COL, runSize - 1 - j);
+ record.putValue(VAL_COL, sortedValue * 1_000_000 + j);
+ record.putValue(NULLABLE_COL, j);
+ records.add(record);
+ }
+ }
+
+ private static List<GenericRow> buildNullTailRecords() {
+ List<GenericRow> records = new ArrayList<>();
+ // Three runs, each with exactly one null tail and two distinct non-null tails, so ordering stays deterministic
+ // (no all-order-by-columns tie) while a null flows through the order-by tail column.
+ for (int g = 0; g < 3; g++) {
+ for (int j = 0; j < 3; j++) {
+ GenericRow record = new GenericRow();
+ record.putValue(SORTED_COL, g);
+ if (j == 0) {
+ record.addNullValueField(TAIL_COL);
+ } else {
+ record.putValue(TAIL_COL, g * 10 + j);
+ }
+ record.putValue(VAL_COL, g * 100 + j);
+ record.putValue(NULLABLE_COL, g * 100 + j);
+ records.add(record);
+ }
+ }
+ return records;
+ }
+
+ private static IndexSegment buildSegment(String segmentName, List<GenericRow> records, boolean nullHandling)
+ throws Exception {
+ SegmentGeneratorConfig segmentGeneratorConfig = new SegmentGeneratorConfig(TABLE_CONFIG, SCHEMA);
+ segmentGeneratorConfig.setTableName(RAW_TABLE_NAME);
+ segmentGeneratorConfig.setSegmentName(segmentName);
+ segmentGeneratorConfig.setDefaultNullHandlingEnabled(nullHandling);
+ segmentGeneratorConfig.setOutDir(TEMP_DIR.getPath());
+
+ SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl();
+ driver.init(segmentGeneratorConfig, new GenericRowRecordReader(records));
+ driver.build();
+
+ return ImmutableSegmentLoader.load(new File(TEMP_DIR, segmentName), ReadMode.mmap);
+ }
+
+ @Test
+ public void testSingleSortedColumnAscending() {
+ assertParity(_segment, "SELECT sortedCol FROM testTable ORDER BY sortedCol", false, 1);
+ }
+
+ @Test
+ public void testSingleSortedColumnDescending() {
+ // Reverse order must be allowed for the forward-scan project to iterate the sorted column descending; otherwise
+ // SelectionPlanNode falls back to the materialized DESC operator and the streaming operator is never built.
+ assertParity(_segment, "SET allowReverseOrder=true; SELECT sortedCol FROM testTable ORDER BY sortedCol DESC", false,
+ 1);
+ }
+
+ @Test
+ public void testSortedPrefixWithUnsortedTail() {
+ // Repeated sorted value + unique tail exercises nextRun(): per-run top-K heap and the one-row lookahead. With the
+ // default LIMIT 10 spanning ~3 runs of 4, the streaming operator emits more than one block.
+ assertParity(_dupSegment, "SELECT sortedCol, tailCol FROM testTable ORDER BY sortedCol, tailCol", false, 2);
+ }
+
+ @Test
+ public void testTwoPhaseSingleSortedColumn() {
+ // Output has a non-order-by column (valCol) -> two-phase fetch. Unique sorted column keeps ordering deterministic.
+ assertParity(_segment, "SELECT valCol, sortedCol FROM testTable ORDER BY sortedCol", false, 1);
+ }
+
+ @Test
+ public void testTwoPhaseWithUnsortedTail() {
+ // Two-phase fetch combined with the run/heap path (sorted prefix + unsorted tail).
+ assertParity(_dupSegment, "SELECT valCol FROM testTable ORDER BY sortedCol, tailCol", false, 1);
+ }
+
+ @Test
+ public void testLimitOffsetNoTail() {
+ // Server retains limit + offset rows; the broker applies the offset later.
+ assertParity(_segment, "SELECT sortedCol FROM testTable ORDER BY sortedCol LIMIT 5 OFFSET 3", false, 1);
+ }
+
+ @Test
+ public void testLimitOffsetWithTail() {
+ assertParity(_dupSegment, "SELECT sortedCol, tailCol FROM testTable ORDER BY sortedCol, tailCol LIMIT 7 OFFSET 5",
+ false, 1);
+ }
+
+ @Test
+ public void testRunSpanningMultipleProjectBlocks() {
+ // One run of 12_000 rows (> one 10k project block); the heap caps at limit + offset while the scan crosses the
+ // block boundary. A single primary value means a single run, hence a single emitted block.
+ StreamingSelectionOrderByOperator operator = assertParity(_largeSegment,
+ "SELECT sortedCol, tailCol FROM testTable ORDER BY sortedCol, tailCol LIMIT 25", false, 1);
+ // The whole run is scanned to find its top-K, which only happens if the forward scan pulled every project block
+ // (proving the scan genuinely crossed the 10k boundary rather than stopping at the first block).
+ assertEquals(operator.getExecutionStatistics().getNumDocsScanned(), NUM_LARGE_RECORDS);
+ }
+
+ @Test
+ public void testRunBoundaryAcrossProjectBlocks() {
+ // First run (10_500 rows) is larger than one project block, and the limit (10_600) exceeds it, so the operator
+ // crosses a project-block boundary mid-run AND carries _pendingRow across the run boundary to emit a second block.
+ StreamingSelectionOrderByOperator operator = assertParity(_multiRunSegment,
+ "SELECT sortedCol, tailCol FROM testTable ORDER BY sortedCol, tailCol LIMIT 10600", false, 2);
+ assertEquals(operator.getExecutionStatistics().getNumDocsScanned(), NUM_MULTI_RUN_RECORDS);
+ }
+
+ @Test
+ public void testNullHandlingEnabled() {
+ assertParity(_dupSegment, "SELECT nullableCol, sortedCol, tailCol FROM testTable ORDER BY sortedCol, tailCol", true,
+ 1);
+ }
+
+ @Test
+ public void testNullHandlingDisabled() {
+ assertParity(_dupSegment, "SELECT nullableCol, sortedCol, tailCol FROM testTable ORDER BY sortedCol, tailCol",
+ false, 1);
+ }
+
+ @Test
+ public void testNullInOrderByColumnWithNullHandling() {
+ // Drives a null through the order-by tail column (not just the carried non-order-by column), exercising the
+ // streaming operator's null-aware comparator and phase-1 null materialization.
+ String query = "SELECT tailCol, sortedCol, valCol FROM testTable ORDER BY sortedCol, tailCol";
+ assertParity(_nullTailSegment, query, true, 1);
+ // Absolute anchor: with null handling on, the order-by tail column (index 1 after extractExpressions) actually
+ // carries nulls through to the output. Guards against both operators substituting a default value identically.
+ List<Object[]> rows = collectStreamingRows(_nullTailSegment, query, true);
+ assertTrue(rows.stream().anyMatch(row -> row[1] == null), "Expected a null in the order-by tail column");
+ }
+
+ @Test
+ public void testNullInOrderByColumnWithoutNullHandling() {
+ String query = "SELECT tailCol, sortedCol, valCol FROM testTable ORDER BY sortedCol, tailCol";
+ assertParity(_nullTailSegment, query, false, 1);
+ // With null handling off the null reads back as the column's default value, so no output cell is null.
+ List<Object[]> rows = collectStreamingRows(_nullTailSegment, query, false);
+ assertTrue(rows.stream().noneMatch(row -> row[1] == null),
+ "Expected no nulls in the output when null handling is disabled");
+ }
+
+ /// Runs {@code query} twice over {@code segment} - once with the streaming hint on, once off - and asserts the
+ /// concatenated streaming blocks equal the materialized operator's single block, cell by cell. Returns the (now
+ /// exhausted) streaming operator so callers can make extra assertions on its execution statistics.
+ ///
+ /// @param expectedMinBlocks the minimum number of non-null blocks the streaming operator must emit. Tail-mode
+ /// cases that span multiple runs pass {@code >= 2} to prove the output is genuinely streamed; cases whose
+ /// result fits in a single trimmed block pass {@code 1}.
+ private StreamingSelectionOrderByOperator assertParity(IndexSegment segment, @Language("sql") String query,
+ boolean nullHandling, int expectedMinBlocks) {
+ // Streaming path.
+ QueryContext streamingContext = QueryContextConverterUtils.getQueryContext(query);
+ streamingContext.setNullHandlingEnabled(nullHandling);
+ streamingContext.setSortedSelectionMergeEnabled(true);
+ Operator<SelectionResultsBlock> streamingOperator =
+ new SelectionPlanNode(new SegmentContext(segment), streamingContext).run();
+ assertTrue(streamingOperator instanceof StreamingSelectionOrderByOperator,
+ "Expected the streaming operator to be built, got: " + streamingOperator.getClass().getSimpleName());
+
+ List<Object[]> streamingRows = new ArrayList<>();
+ DataSchema streamingSchema = null;
+ int numBlocks = 0;
+ SelectionResultsBlock block;
+ while ((block = streamingOperator.nextBlock()) != null) {
+ numBlocks++;
+ if (streamingSchema == null) {
+ streamingSchema = block.getDataSchema();
+ }
+ streamingRows.addAll(block.getRows());
+ }
+ assertTrue(numBlocks >= expectedMinBlocks,
+ "Expected at least " + expectedMinBlocks + " streaming block(s), got: " + numBlocks);
+
+ // Materialized path (hint off).
+ QueryContext materializedContext = QueryContextConverterUtils.getQueryContext(query);
+ materializedContext.setNullHandlingEnabled(nullHandling);
+ Operator<SelectionResultsBlock> materializedOperator =
+ new SelectionPlanNode(new SegmentContext(segment), materializedContext).run();
+ assertFalse(materializedOperator instanceof StreamingSelectionOrderByOperator,
+ "Materialized baseline must not be the streaming operator");
+ SelectionResultsBlock materializedBlock = materializedOperator.nextBlock();
+ assertNotNull(materializedBlock);
+ List<Object[]> expectedRows = materializedBlock.getRows();
+
+ assertEquals(streamingSchema, materializedBlock.getDataSchema(), "Schema mismatch for query: " + query);
+ assertEquals(streamingRows.size(), expectedRows.size(), "Row count mismatch for query: " + query);
+ for (int i = 0; i < expectedRows.size(); i++) {
+ assertEquals(streamingRows.get(i), expectedRows.get(i), "Row " + i + " mismatch for query: " + query);
+ }
+ return (StreamingSelectionOrderByOperator) streamingOperator;
+ }
+
+ /// Drains the streaming operator for {@code query} and returns all rows concatenated in emission order.
+ private List<Object[]> collectStreamingRows(IndexSegment segment, @Language("sql") String query,
+ boolean nullHandling) {
+ QueryContext queryContext = QueryContextConverterUtils.getQueryContext(query);
+ queryContext.setNullHandlingEnabled(nullHandling);
+ queryContext.setSortedSelectionMergeEnabled(true);
+ Operator<SelectionResultsBlock> operator = new SelectionPlanNode(new SegmentContext(segment), queryContext).run();
+ List<Object[]> rows = new ArrayList<>();
+ SelectionResultsBlock block;
+ while ((block = operator.nextBlock()) != null) {
+ rows.addAll(block.getRows());
+ }
+ return rows;
+ }
+
+ @AfterClass
+ public void tearDown()
+ throws IOException {
+ for (IndexSegment segment : new IndexSegment[]{_segment, _dupSegment, _largeSegment, _multiRunSegment,
+ _nullTailSegment}) {
+ if (segment != null) {
+ segment.destroy();
+ }
+ }
+ FileUtils.deleteDirectory(TEMP_DIR);
+ }
+}
diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
index 7c5db4ea1f2..0af0e797710 100644
--- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
+++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
@@ -630,6 +630,11 @@ public class CommonConstants {
public static final String CONFIG_OF_MSE_STREAMING_DISTINCT_FLUSH_THRESHOLD =
"pinot.broker.mse.streaming.distinct.flush.threshold";
public static final int DEFAULT_MSE_STREAMING_DISTINCT_FLUSH_THRESHOLD = -1;
+
+ /// Default output block size (rows) for the streaming selection ORDER BY combine
+ /// ({@link Request.QueryOptionKey#SORTED_SELECTION_MERGE_BLOCK_SIZE}).
+ public static final int DEFAULT_SORTED_SELECTION_MERGE_BLOCK_SIZE = 10_000;
+
// Whether to infer partition hint by default or not.
// This value can always be overridden by INFER_PARTITION_HINT query option
public static final String CONFIG_OF_INFER_PARTITION_HINT = "pinot.broker.multistage.infer.partition.hint";
@@ -925,6 +930,11 @@ public class CommonConstants {
/// there is no such stage and the client would observe duplicate rows across flush windows.
public static final String STREAMING_DISTINCT_FLUSH_THRESHOLD = "streamingDistinctFlushThreshold";
+ /// Opt-in: use the streaming k-way-merge selection ORDER BY combine over sorted segments.
+ public static final String SORTED_SELECTION_MERGE_ENABLED = "sortedSelectionMergeEnabled";
+ /// Output block size (rows) for the streaming selection ORDER BY combine.
+ public static final String SORTED_SELECTION_MERGE_BLOCK_SIZE = "sortedSelectionMergeBlockSize";
+
public static final String NUM_REPLICA_GROUPS_TO_QUERY = "numReplicaGroupsToQuery";
public static final String ORDERED_PREFERRED_POOLS = "orderedPreferredPools";
public static final String USE_FIXED_REPLICA = "useFixedReplica";
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
(pinot) 05/07: Keep frontier pruning for segments flagged non-null.
Posted by xi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
xiangfu0 pushed a commit to branch xiangfu0/codex1/sorted-exchange-leaf
in repository https://gitbox.apache.org/repos/asf/pinot.git
commit 7efe592da816122a337106b02b4f010617357f45
Author: rohity <ro...@uber.com>
AuthorDate: Mon Sep 28 04:20:12 2026 +0000
Keep frontier pruning for segments flagged non-null.
A query-level null-handling switch turned pruning off for every segment. The non-null metadata flag is already loaded with the segment, so flagged columns keep their min/max and unflagged ones always activate. AUTO uses the same predicate and no longer selects the streaming merge over segments the leaf will not stream. A consuming segment has no column metadata map, so it counts as unflagged rather than failing the query.
---
.../StreamingSelectionOrderByCombineOperator.java | 21 +-
.../apache/pinot/core/plan/SelectionPlanNode.java | 20 +-
.../core/plan/maker/InstancePlanMakerImplV2.java | 30 +--
...reamingSelectionOrderByCombineOperatorTest.java | 212 +++++++++++++++++++--
4 files changed, 250 insertions(+), 33 deletions(-)
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperator.java
index 5597a06bcfc..7ecf0a8cb9f 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperator.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperator.java
@@ -38,6 +38,7 @@ 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;
@@ -73,8 +74,11 @@ import org.slf4j.LoggerFactory;
/// 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.
-/// 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.
+/// 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`.
@@ -165,7 +169,7 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
String firstOrderByColumn =
firstOrderByExpressionContext.getType() == ExpressionContext.Type.IDENTIFIER
? firstOrderByExpressionContext.getIdentifier() : null;
- _pruningEnabled = firstOrderByColumn != null && !queryContext.isNullHandlingEnabled();
+ _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).
@@ -185,7 +189,16 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
DataSourceMetadata metadata =
operator.getIndexSegment().getDataSource(firstOrderByColumn, queryContext.getSchema())
.getDataSourceMetadata();
- _sortedCursors[i] = new SegmentCursor(operator, metadata.getMinValue(), metadata.getMaxValue());
+ Comparable minValue = metadata.getMinValue();
+ Comparable maxValue = metadata.getMaxValue();
+ // 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++;
}
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/plan/SelectionPlanNode.java b/pinot-core/src/main/java/org/apache/pinot/core/plan/SelectionPlanNode.java
index f2a39b14777..c022bd76173 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/plan/SelectionPlanNode.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/plan/SelectionPlanNode.java
@@ -20,6 +20,7 @@ package org.apache.pinot.core.plan;
import java.util.ArrayList;
import java.util.List;
+import java.util.Map;
import org.apache.pinot.common.request.context.ExpressionContext;
import org.apache.pinot.common.request.context.OrderByExpressionContext;
import org.apache.pinot.common.utils.config.QueryOptionsUtils;
@@ -35,6 +36,7 @@ import org.apache.pinot.core.operator.query.SelectionPartiallyOrderedByLinearOpe
import org.apache.pinot.core.operator.query.StreamingSelectionOrderByOperator;
import org.apache.pinot.core.query.request.context.QueryContext;
import org.apache.pinot.core.query.selection.SelectionOperatorUtils;
+import org.apache.pinot.segment.spi.ColumnMetadata;
import org.apache.pinot.segment.spi.IndexSegment;
import org.apache.pinot.segment.spi.SegmentContext;
import org.apache.pinot.segment.spi.datasource.DataSource;
@@ -231,7 +233,9 @@ public class SelectionPlanNode implements PlanNode {
/// This is [#isColumnPhysicallySorted] plus the null caveat: once null handling is on, the physical order is not
/// the order the query asks for, so a column carrying nulls is not usable as sorted. Reading the null bitmap
/// touches the segment's mapped buffer, so this must only be called once the segment has been acquired -- which is
- /// why [AcquireReleaseColumnsSegmentPlanNode] defers the whole plan build until after `acquire()`.
+ /// why [AcquireReleaseColumnsSegmentPlanNode] defers the whole plan build until after `acquire()`. Callers that
+ /// run before acquire use [#isColumnFlaggedNonNull] instead. That flag is one-sided (`true` means no nulls; `false`
+ /// means nulls or unknown) and is stricter than the bitmap check below.
public static boolean isColumnSorted(IndexSegment segment, QueryContext queryContext, String column) {
DataSource dataSource = segment.getDataSource(column, queryContext.getSchema());
// If there are null values, we cannot trust DataSourceMetadata.isSorted
@@ -253,4 +257,18 @@ public class SelectionPlanNode implements PlanNode {
public static boolean isColumnPhysicallySorted(IndexSegment segment, QueryContext queryContext, String column) {
return segment.getDataSource(column, queryContext.getSchema()).getDataSourceMetadata().isSorted();
}
+
+ /// Returns whether segment metadata flags `column` as holding no nulls ([ColumnMetadata#isNonNull]).
+ ///
+ /// Reads segment metadata only, so it needs no segment acquire. `false` means nulls or unknown. A consuming segment
+ /// has no column metadata map (it is `null`, so
+ /// [org.apache.pinot.segment.spi.SegmentMetadata#getColumnMetadataFor] would throw) and is never flagged.
+ public static boolean isColumnFlaggedNonNull(IndexSegment segment, String column) {
+ Map<String, ColumnMetadata> columnMetadataMap = segment.getSegmentMetadata().getColumnMetadataMap();
+ if (columnMetadataMap == null) {
+ return false;
+ }
+ ColumnMetadata columnMetadata = columnMetadataMap.get(column);
+ return columnMetadata != null && columnMetadata.isNonNull();
+ }
}
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java b/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java
index f00ffb4aa18..6b7aa155c98 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java
@@ -505,28 +505,30 @@ public class InstancePlanMakerImplV2 implements PlanMaker {
//
// TODO: allowReverseOrder=true still does not guarantee streaming leaves. getSortedByProject() swallows the
// failure when a docId set cannot be reversed and returns the unreversed operator, so the leaf gate then
- // declines and AUTO has again resolved ON over materialized children. Same residual shape as the null
- // handling case below; a shared "can this segment actually stream in the query's direction" predicate
- // would close both. See https://github.com/apache/pinot/pull/19120#discussion_r3871713989
+ // declines and AUTO has resolved ON over materialized children.
+ // See https://github.com/apache/pinot/pull/19120#discussion_r3871713989
LOGGER.debug("Not using {} for a DESC leading ORDER BY expression without {}: {}",
QueryOptionKey.SORTED_SELECTION_MERGE_MODE, QueryOptionKey.ALLOW_REVERSE_ORDER, firstOrderByExpression);
return false;
}
String column = firstOrderByExpression.getIdentifier();
+ boolean nullHandlingEnabled = queryContext.isNullHandlingEnabled();
int numSorted = 0;
for (SegmentContext segmentContext : segmentContexts) {
- // TODO: This deliberately uses the physical-sortedness predicate rather than the full
- // SelectionPlanNode.isColumnSorted(), which additionally rejects a column carrying nulls when null
- // handling is on. That null check reads the segment's mapped buffer via the null value vector, and this
- // runs before any segment is acquired -- AcquireReleaseColumnsSegmentPlanNode defers the whole plan build
- // precisely so that no planner touches a buffer pre-acquire. Consequence: with null handling on and a
- // null-bearing leading column, AUTO can resolve to ON while the leaf gate then refuses to stream, giving
- // a k-way merge over fully materialized children -- correct, but slower than the MinMax combine. Handle
- // properly (an acquire-free nullability signal, or resolving AUTO per segment at acquire time) as a
- // follow-up. See https://github.com/apache/pinot/pull/19120#discussion_r3871713975
- if (SelectionPlanNode.isColumnPhysicallySorted(segmentContext.getIndexSegment(), queryContext, column)) {
- numSorted++;
+ // Physical sortedness is metadata. SelectionPlanNode.isColumnSorted() also rejects a null-bearing column, but
+ // that reads the null value vector, a mapped buffer, and this runs before any segment is acquired. Use
+ // SelectionPlanNode.isColumnFlaggedNonNull() instead. It is one-sided (false means nulls or unknown) and
+ // stricter than the leaf, which can see an empty bitmap after acquire: an old segment with no flag resolves OFF
+ // rather than streaming over a materialized child.
+ // See https://github.com/apache/pinot/pull/19120#discussion_r3871713975
+ IndexSegment segment = segmentContext.getIndexSegment();
+ if (!SelectionPlanNode.isColumnPhysicallySorted(segment, queryContext, column)) {
+ continue;
}
+ if (nullHandlingEnabled && !SelectionPlanNode.isColumnFlaggedNonNull(segment, column)) {
+ continue;
+ }
+ numSorted++;
}
double sortedRatio = (double) numSorted / numSegments;
double minSortedRatio = queryContext.getSortedSelectionMergeAutoMinSortedRatio();
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java
index 606705708cc..2532d4aacc7 100644
--- a/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java
+++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java
@@ -24,6 +24,7 @@ import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.Map;
+import java.util.Set;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.stream.Collectors;
@@ -43,6 +44,7 @@ import org.apache.pinot.core.operator.query.StreamingSelectionOrderByOperator;
import org.apache.pinot.core.plan.CombinePlanNode;
import org.apache.pinot.core.plan.ExplainInfo;
import org.apache.pinot.core.plan.PlanNode;
+import org.apache.pinot.core.plan.SelectionPlanNode;
import org.apache.pinot.core.plan.maker.InstancePlanMakerImplV2;
import org.apache.pinot.core.plan.maker.PlanMaker;
import org.apache.pinot.core.query.executor.ResultsBlockStreamer;
@@ -51,8 +53,11 @@ import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUt
import org.apache.pinot.core.query.utils.OrderByComparatorFactory;
import org.apache.pinot.core.util.QueryMultiThreadingUtils;
import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader;
+import org.apache.pinot.segment.local.indexsegment.mutable.MutableSegmentImpl;
+import org.apache.pinot.segment.local.indexsegment.mutable.MutableSegmentImplTestUtils;
import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl;
import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader;
+import org.apache.pinot.segment.spi.ColumnMetadata;
import org.apache.pinot.segment.spi.IndexSegment;
import org.apache.pinot.segment.spi.SegmentContext;
import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig;
@@ -161,6 +166,14 @@ public class StreamingSelectionOrderByCombineOperatorTest {
/// Two segments whose rows tie on every order-by expression; see [#buildTiedSortedRecords].
private List<IndexSegment> _tiedSegments;
+ /// Three disjoint ranges built with null handling and no nulls in `sortedCol`, so the column is flagged non-null.
+ /// Ranges are `[0, 100)`, `[1000, 1100)`, `[2000, 2100)`.
+ private List<IndexSegment> _nonNullDisjointSegments;
+ /// The middle of those ranges, with nulls in `sortedCol`. Stored nulls are `Integer.MIN_VALUE`, so the metadata min
+ /// is not an ASC bound; the metadata max is the highest non-null, which a DESC top-5 over the high range would prune
+ /// if the flag were ignored.
+ private IndexSegment _nullBearingMidSegment;
+
@BeforeClass
public void setUp()
throws Exception {
@@ -210,6 +223,14 @@ public class StreamingSelectionOrderByCombineOperatorTest {
_interleavedSegments = new ArrayList<>(2);
_interleavedSegments.add(buildSegment(SORTED_TABLE_CONFIG, "sparse_0", buildSparseSortedRecords(), false));
_interleavedSegments.add(buildSegment(SORTED_TABLE_CONFIG, "denseGap_0", buildGapFillingSortedRecords(), false));
+
+ _nonNullDisjointSegments = new ArrayList<>(3);
+ for (int i = 0; i < 3; i++) {
+ _nonNullDisjointSegments.add(
+ buildSegment(SORTED_TABLE_CONFIG, "nonNullDisjoint_" + i, buildDisjointSortedRecords(i), true));
+ }
+ _nullBearingMidSegment = buildSegment(SORTED_TABLE_CONFIG, "nullBearingMid_0",
+ buildNullBearingDisjointRecords(1), true);
}
/// Leading order-by values `0, 1000, 2000, ...`: gaps wide enough for another segment's entire range to sit between
@@ -266,6 +287,17 @@ public class StreamingSelectionOrderByCombineOperatorTest {
return records;
}
+ /// Like [#buildDisjointSortedRecords] but the leading order-by column itself is null in the first few rows.
+ private static List<GenericRow> buildNullBearingDisjointRecords(int index) {
+ List<GenericRow> records = buildDisjointSortedRecords(index);
+ for (int i = 0; i < 3; i++) {
+ GenericRow record = records.get(i);
+ record.putValue(SORTED_COL, null);
+ record.addNullValueField(SORTED_COL);
+ }
+ return records;
+ }
+
/// Like [#buildOverlappingSortedRecords] but the leading order-by column itself is null in the first few rows.
private static List<GenericRow> buildNullBearingSortedRecords(int index) {
List<GenericRow> records = buildOverlappingSortedRecords(index);
@@ -365,6 +397,32 @@ public class StreamingSelectionOrderByCombineOperatorTest {
return ImmutableSegmentLoader.load(new File(TEMP_DIR, segmentName), ReadMode.mmap);
}
+ /// [ColumnMetadata#isNonNull] is the acquire-free signal pruning and AUTO use. These fixtures have to carry it the
+ /// way [org.apache.pinot.segment.local.segment.creator.impl.BaseSegmentCreator] writes it, or the tests below are
+ /// vacuous.
+ @Test
+ public void testNonNullFlagMatchesHowSegmentsWereBuilt() {
+ for (IndexSegment segment : _sortedSegments) {
+ assertTrue(isColumnNonNull(segment, SORTED_COL), segment.getSegmentName());
+ assertFalse(isColumnNonNull(segment, NULLABLE_COL), segment.getSegmentName());
+ }
+ for (IndexSegment segment : _disjointSegments) {
+ assertFalse(isColumnNonNull(segment, SORTED_COL), segment.getSegmentName());
+ }
+ for (IndexSegment segment : _nullBearingSortedSegments) {
+ assertFalse(isColumnNonNull(segment, SORTED_COL), segment.getSegmentName());
+ }
+ for (IndexSegment segment : _nonNullDisjointSegments) {
+ assertTrue(isColumnNonNull(segment, SORTED_COL), segment.getSegmentName());
+ }
+ assertFalse(isColumnNonNull(_nullBearingMidSegment, SORTED_COL));
+ }
+
+ private static boolean isColumnNonNull(IndexSegment segment, String column) {
+ ColumnMetadata metadata = segment.getSegmentMetadata().getColumnMetadataFor(column);
+ return metadata != null && metadata.isNonNull();
+ }
+
@Test
public void testAscendingParity() {
assertParity(_sortedSegments, "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50", false);
@@ -587,7 +645,7 @@ public class StreamingSelectionOrderByCombineOperatorTest {
@Test
public void testNullHandlingEnabledParity() {
- // Null handling on disables min/max pruning (the combine activates every segment); nullableCol carries real nulls.
+ // nullableCol carries real nulls. sortedCol is flagged non-null, so frontier pruning stays on for this order-by.
assertParity(_sortedSegments,
"SELECT nullableCol, sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50", true);
}
@@ -635,6 +693,126 @@ public class StreamingSelectionOrderByCombineOperatorTest {
assertEquals(result._numSegmentsMatched, 1, "numSegmentsMatched must exclude never-activated segments");
}
+ /// Null handling on, but every segment's leading column is flagged non-null, so min/max pruning stays exact.
+ @Test
+ public void testNullHandlingPrunesFlaggedNonNullSegments() {
+ @Language("sql") String query =
+ "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 5";
+ assertParity(_nonNullDisjointSegments, query, true);
+ Result result = run(_nonNullDisjointSegments, query, true, true, true, 0);
+ assertEquals(result._rows.size(), 5);
+ assertSegmentsScanned(result, 1, _nonNullDisjointSegments.size());
+ }
+
+ /// These segments were built with null handling off, so `sortedCol` is unflagged even though it holds no nulls.
+ /// With null handling on the query, every segment must activate: an unknown flag is not proof of no nulls.
+ @Test
+ public void testUnflaggedSegmentsAlwaysActivateUnderNullHandling() {
+ @Language("sql") String query =
+ "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 5";
+ Result result = run(_disjointSegments, query, true, true, true, 0);
+ assertEquals(result._rows.size(), 5);
+ assertSegmentsScanned(result, NUM_SEGMENTS, NUM_SEGMENTS);
+ }
+
+ /// Flagged low and high ranges around one null-bearing middle segment. ASC top-5 lives entirely in the low range;
+ /// the high range prunes, and the null-bearing segment still activates. On DESC a null frontier head stops further
+ /// pruning, so the skipped-segment claim is made by acquiring the null-bearing segment rather than by a scan count.
+ @Test
+ public void testNullBearingSegmentForceActivatesAlongsideFlaggedOnes() {
+ List<IndexSegment> segments = nullPruningMix();
+ @Language("sql") String query =
+ "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 5";
+ assertParity(segments, query, true);
+ Result result = run(segments, query, true, true, true, 0);
+ assertEquals(result._rows.size(), 5);
+ for (int i = 0; i < 5; i++) {
+ assertEquals((int) result._rows.get(i)[0], i);
+ }
+ assertSegmentsScanned(result, 2, segments.size());
+ // The ASC count above does not prove this segment was force-activated: its stored null is Integer.MIN_VALUE, so
+ // the metadata min would activate it anyway. On DESC the metadata max is the highest non-null and sits below the
+ // high range, so keeping that max would never acquire it. Its nulls sort first, which is what parity checks.
+ @Language("sql") String desc =
+ "SET allowReverseOrder=true; SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol DESC, valCol DESC "
+ + "LIMIT 5";
+ assertParity(segments, desc, true);
+ QueryContext descContext = hintedContext(desc);
+ descContext.setNullHandlingEnabled(true);
+ List<InstrumentedSegmentOperator> children = instrument(segments, descContext);
+ drain(combineOver(children, descContext));
+ assertTrue(children.get(1)._numAcquires > 0,
+ "The null-bearing segment must be acquired; its metadata max would prune it");
+ }
+
+ /// DESC with `allowReverseOrder`. The low segment was built with null handling off, so it is unflagged and its max
+ /// sits below the high range: trusting that max would prune it, and with it the reverse bitmap built on the first
+ /// block. The middle flagged segment stays past the frontier and is never scanned. A null-bearing segment cannot
+ /// stand in here, because a null frontier head refuses every later prune.
+ @Test
+ public void testDescReverseScanOnlyForceActivatesUnflaggedSegments() {
+ List<IndexSegment> segments = List.of(_disjointSegments.get(0), _nonNullDisjointSegments.get(1),
+ _nonNullDisjointSegments.get(2));
+ @Language("sql") String query =
+ "SET allowReverseOrder=true; SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol DESC, valCol DESC "
+ + "LIMIT 5";
+ assertParity(segments, query, true);
+ Result result = run(segments, query, true, true, true, 0);
+ assertEquals(result._rows.size(), 5);
+ int top = 2 * 1000 + NUM_RECORDS_PER_SEGMENT - 1;
+ for (int i = 0; i < 5; i++) {
+ assertEquals((int) result._rows.get(i)[0], top - i);
+ }
+ assertSegmentsScanned(result, 2, segments.size());
+ }
+
+ /// A consuming segment has no column metadata map, so reading its non-null flag through
+ /// `SegmentMetadata.getColumnMetadataFor` throws. It tracks min/max as it ingests, though, so if it were treated as
+ /// flagged its bounds would prune it here just like the high segment. It has to count as unflagged and activate.
+ @Test
+ public void testConsumingSegmentAlwaysActivatesUnderNullHandling()
+ throws Exception {
+ MutableSegmentImpl consumingSegment =
+ MutableSegmentImplTestUtils.createMutableSegmentImpl(SCHEMA, Set.of(), Set.of(), Set.of(), false, true);
+ try {
+ for (GenericRow record : buildDisjointSortedRecords(1)) {
+ consumingSegment.index(record, null);
+ }
+ assertNull(consumingSegment.getSegmentMetadata().getColumnMetadataMap(),
+ "The fixture must reproduce the consuming segment's missing column metadata, or this test is vacuous");
+ assertFalse(SelectionPlanNode.isColumnFlaggedNonNull(consumingSegment, SORTED_COL));
+
+ List<IndexSegment> segments =
+ List.of(_nonNullDisjointSegments.get(0), consumingSegment, _nonNullDisjointSegments.get(2));
+ @Language("sql") String query = "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 5";
+ assertParity(segments, query, true);
+ Result result = run(segments, query, true, true, true, 0);
+ assertEquals(result._rows.size(), 5);
+ for (int i = 0; i < 5; i++) {
+ assertEquals((int) result._rows.get(i)[0], i);
+ }
+ assertSegmentsScanned(result, 2, segments.size());
+ } finally {
+ consumingSegment.destroy();
+ }
+ }
+
+ /// Low flagged, null-bearing middle, high flagged. The middle segment is the only one that must force-activate.
+ private List<IndexSegment> nullPruningMix() {
+ return List.of(_nonNullDisjointSegments.get(0), _nullBearingMidSegment, _nonNullDisjointSegments.get(2));
+ }
+
+ /// `numSegmentsMatched` is the count of children that scanned a doc, which is the cursors the merge activated. A
+ /// pruned child scans nothing, which is also what keeps it from building a DESC reverse bitmap. Doc count is only
+ /// an upper bound: an activated streaming child may stop after the first rows of its block.
+ private static void assertSegmentsScanned(Result result, int numScanned, int numSegments) {
+ assertEquals(result._numSegmentsProcessed, numSegments, "numSegmentsProcessed counts every segment");
+ assertEquals(result._numSegmentsMatched, numScanned,
+ "numSegmentsMatched counts segments that scanned a doc; scanned docs: " + result._numDocsScanned);
+ assertTrue(result._numDocsScanned > 0 && result._numDocsScanned <= (long) numScanned * NUM_RECORDS_PER_SEGMENT,
+ "Scanned " + result._numDocsScanned + " docs across " + numScanned + " segments");
+ }
+
/// A cursor activated *late* can hold a smaller row than the retained leader. Two things must be right for it to
/// land in the correct place: the leader is re-compared after `activateEligibleCursors()` rather than before, and
/// the pruning frontier is the row about to be emitted rather than the last one emitted. Either mistake emits the
@@ -1077,9 +1255,10 @@ public class StreamingSelectionOrderByCombineOperatorTest {
public void testExplainAttributesReflectTheQueryShape() {
assertTrue(explainAttributes(_sortedSegments, "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol LIMIT 5",
false).get("deferTiedCursors").getBool(), "A single order-by expression defers ties");
- assertFalse(explainAttributes(_sortedSegments,
+ assertTrue(explainAttributes(_sortedSegments,
"SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 5", true)
- .get("frontierPruning").getBool(), "Null handling disables frontier pruning");
+ .get("frontierPruning").getBool(),
+ "Null handling no longer disables frontier pruning when the leading order-by is a column");
assertEquals(explainAttributes(_mixedSegments,
"SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 5", false)
.get("numSortedSegments").getLong(), 2);
@@ -1202,14 +1381,11 @@ public class StreamingSelectionOrderByCombineOperatorTest {
}
@Test
- public void testAutoIgnoresNullsInTheLeadingColumn() {
- // Pins the known gap documented by the TODO in InstancePlanMakerImplV2#isSortedEnoughForStreamingMerge, so it is
- // enforced by CI rather than only described. AUTO resolves from physical sortedness alone, because the null check
- // that SelectionPlanNode#isColumnSorted adds reads the segment's mapped buffer and AUTO runs before any acquire.
- // So a null-bearing sorted column resolves to ON with null handling either off or on -- but with it on, the leaf
- // gate then refuses to stream, and the merge runs over fully materialized children. Correct, just not fast.
- //
- // Tighten this to OFF for the null-handling case when the follow-up lands.
+ public void testAutoRequiresNonNullLeadingColumnWhenNullHandlingIsOn() {
+ // AUTO runs before any segment is acquired, so it must not read the null value vector. It uses
+ // ColumnMetadata.isNonNull() instead, which is one-sided: false means nulls or unknown. That is stricter than
+ // SelectionPlanNode.isColumnSorted(), which still reads the bitmap after acquire. An old segment with an empty
+ // bitmap and no flag therefore resolves OFF rather than streaming over a materialized child.
// See https://github.com/apache/pinot/pull/19120#discussion_r3871713975
String query = "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 50";
@@ -1221,8 +1397,8 @@ public class StreamingSelectionOrderByCombineOperatorTest {
QueryContext nullHandlingOn = modeContext(query, SortedSelectionMergeMode.AUTO);
nullHandlingOn.setNullHandlingEnabled(true);
makeStreamingInstancePlan(_nullBearingSortedSegments, nullHandlingOn);
- assertEquals(nullHandlingOn.getSortedSelectionMergeMode(), SortedSelectionMergeMode.ON,
- "AUTO currently reads physical sortedness only, so null handling does not change the decision");
+ assertEquals(nullHandlingOn.getSortedSelectionMergeMode(), SortedSelectionMergeMode.OFF,
+ "With null handling on an unflagged leading column must not count as sorted");
// The leaf, which runs post-acquire and can afford the null check, still refuses to stream these segments.
List<SegmentContext> segmentContexts = new ArrayList<>(1);
segmentContexts.add(new SegmentContext(_nullBearingSortedSegments.get(0)));
@@ -1230,6 +1406,12 @@ public class StreamingSelectionOrderByCombineOperatorTest {
assertFalse(leafOperator instanceof StreamingSelectionOrderByOperator,
"A null-bearing leading column must not produce a streaming leaf under null handling, got: "
+ leafOperator.getClass().getSimpleName());
+
+ QueryContext flagged = modeContext(query, SortedSelectionMergeMode.AUTO);
+ flagged.setNullHandlingEnabled(true);
+ makeStreamingInstancePlan(_sortedSegments, flagged);
+ assertEquals(flagged.getSortedSelectionMergeMode(), SortedSelectionMergeMode.ON,
+ "A leading column flagged non-null stays sorted under null handling");
}
@Test
@@ -1489,11 +1671,13 @@ public class StreamingSelectionOrderByCombineOperatorTest {
throws IOException {
EXECUTOR.shutdownNow();
for (List<IndexSegment> segments : List.of(_sortedSegments, _disjointSegments, _mixedSegments, _lowCardSegments,
- _interleavedSegments, _tiedSegments)) {
+ _interleavedSegments, _tiedSegments, _nullBearingSortedSegments, _unsortedSegments,
+ _nonNullDisjointSegments)) {
for (IndexSegment segment : segments) {
segment.destroy();
}
}
+ _nullBearingMidSegment.destroy();
FileUtils.deleteDirectory(TEMP_DIR);
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
(pinot) 07/07: Share sorted selection setup and trim duplicate coverage
Posted by xi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
xiangfu0 pushed a commit to branch xiangfu0/codex1/sorted-exchange-leaf
in repository https://gitbox.apache.org/repos/asf/pinot.git
commit eecbbd1c04489681085b64828c0118c37c958121
Author: Xiang Fu <xi...@apache.org>
AuthorDate: Wed Sep 30 20:55:15 2026 +0530
Share sorted selection setup and trim duplicate coverage
Use one project-block and schema setup path, releasing no-tail block references after materialization. Keep fallback, ties, nulls and DESC coverage while replacing repeated empty-block fixtures with a small control over real segment reads.
---
.../StreamingSelectionOrderByCombineOperator.java | 32 +-
.../query/StreamingSelectionOrderByOperator.java | 121 +++-----
...reamingSelectionOrderByCombineOperatorTest.java | 74 +----
.../StreamingSelectionOrderByOperatorTest.java | 342 ++-------------------
4 files changed, 104 insertions(+), 465 deletions(-)
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperator.java
index 7ecf0a8cb9f..34d90e2a02d 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperator.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperator.java
@@ -217,27 +217,17 @@ public class StreamingSelectionOrderByCombineOperator extends BaseStreamingCombi
/// 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);
- });
- }
+ Arrays.sort(_sortedCursors, (o1, o2) -> {
+ Comparable bound1 = _asc ? o1._minValue : o1._maxValue;
+ Comparable bound2 = _asc ? o2._minValue : o2._maxValue;
+ if (bound1 == null) {
+ return bound2 == null ? 0 : -1;
+ }
+ if (bound2 == null) {
+ return 1;
+ }
+ return _asc ? bound1.compareTo(bound2) : bound2.compareTo(bound1);
+ });
}
@Override
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperator.java
index 58f5f37ef7b..3010b9606f6 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperator.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperator.java
@@ -34,6 +34,7 @@ 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.common.utils.HashUtil;
import org.apache.pinot.core.common.BlockValSet;
import org.apache.pinot.core.common.Operator;
@@ -196,7 +197,7 @@ public class StreamingSelectionOrderByOperator extends BaseOperator<SelectionRes
_phase2NumColumns = 0;
_phase2DataSourceMap = null;
// Single-phase: all output expressions are order-by expressions, so their types are known up front.
- _dataSchema = buildSinglePhaseDataSchema();
+ _dataSchema = buildDataSchema(_projectOperator::getResultColumnContext);
}
_numPhase1Columns = _phase1Expressions.size();
}
@@ -216,7 +217,7 @@ public class StreamingSelectionOrderByOperator extends BaseOperator<SelectionRes
// 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);
+ _dataSchema = buildDataSchema(this::resolveResultColumnContext);
}
assert _dataSchema != null;
return new SelectionResultsBlock(_dataSchema, List.of(), _comparator, _queryContext);
@@ -240,34 +241,20 @@ public class StreamingSelectionOrderByOperator extends BaseOperator<SelectionRes
if (remaining <= 0) {
return null;
}
- ValueBlock valueBlock = nextNonEmptyBlock();
- if (valueBlock == null) {
+ if (!loadNextBlock()) {
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);
+ int numRows = Math.min(_currentNumDocs, remaining);
List<Object[]> rows = new ArrayList<>(numRows);
for (int i = 0; i < numRows; i++) {
- rows.add(materializeRow(blockValueFetcher, docIds, nullBitmaps, i));
+ rows.add(materializeRow(_currentFetcher, _currentDocIds, _currentNullBitmaps, i));
}
+ _currentBlock = null;
+ _currentFetcher = null;
+ _currentDocIds = null;
+ _currentNullBitmaps = null;
_numRowsEmitted += rows.size();
return rows;
}
@@ -319,37 +306,42 @@ public class StreamingSelectionOrderByOperator extends BaseOperator<SelectionRes
/// 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());
+ if ((_currentBlock == null || _currentPos >= _currentNumDocs) && !loadNextBlock()) {
+ return null;
}
int rowId = _currentPos++;
return materializeRow(_currentFetcher, _currentDocIds, _currentNullBitmaps, rowId);
}
+ private boolean loadNextBlock() {
+ if (_projectExhausted) {
+ return false;
+ }
+ _currentBlock = nextNonEmptyBlock();
+ if (_currentBlock == null) {
+ _projectExhausted = true;
+ return false;
+ }
+ 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());
+ return true;
+ }
+
/// 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.
@@ -461,22 +453,11 @@ public class StreamingSelectionOrderByOperator extends BaseOperator<SelectionRes
}
if (_dataSchema == null) {
- _dataSchema = buildTwoPhaseDataSchema(transformOperator::getResultColumnContext);
+ _dataSchema = buildDataSchema(transformOperator::getResultColumnContext);
}
}
}
- private DataSchema buildSinglePhaseDataSchema() {
- String[] columnNames = new String[_numExpressions];
- DataSchema.ColumnDataType[] columnDataTypes = new DataSchema.ColumnDataType[_numExpressions];
- for (int i = 0; i < _numExpressions; i++) {
- columnNames[i] = _expressions.get(i).toString();
- columnDataTypes[i] = DataSchema.ColumnDataType.fromDataType(_orderByColumnContexts[i].getDataType(),
- _orderByColumnContexts[i].isSingleValue());
- }
- return new DataSchema(columnNames, columnDataTypes);
- }
-
/// Resolves a non-order-by expression's result type without building the phase-2 pipeline, for the zero-match case
/// where there is nothing to fetch. This reproduces [TransformOperator#getResultColumnContext] exactly -- that
/// method resolves against its project operator's source column contexts, which for phase 2 are
@@ -494,21 +475,15 @@ public class StreamingSelectionOrderByOperator extends BaseOperator<SelectionRes
TransformFunctionFactory.get(expression, _phase2ColumnContextMap, _queryContext));
}
- private DataSchema buildTwoPhaseDataSchema(Function<ExpressionContext, ColumnContext> resultColumnContexts) {
- int numNonOrderByExpressions = _nonOrderByExpressions.size();
+ private DataSchema buildDataSchema(Function<ExpressionContext, ColumnContext> resultColumnContexts) {
String[] columnNames = new String[_numExpressions];
- DataSchema.ColumnDataType[] columnDataTypes = new DataSchema.ColumnDataType[_numExpressions];
+ ColumnDataType[] columnDataTypes = new ColumnDataType[_numExpressions];
for (int i = 0; i < _numExpressions; i++) {
columnNames[i] = _expressions.get(i).toString();
- }
- for (int i = 0; i < _numOrderByExpressions; i++) {
- columnDataTypes[i] = DataSchema.ColumnDataType.fromDataType(_orderByColumnContexts[i].getDataType(),
- _orderByColumnContexts[i].isSingleValue());
- }
- for (int i = 0; i < numNonOrderByExpressions; i++) {
- ColumnContext columnContext = resultColumnContexts.apply(_nonOrderByExpressions.get(i));
- columnDataTypes[_numOrderByExpressions + i] =
- DataSchema.ColumnDataType.fromDataType(columnContext.getDataType(), columnContext.isSingleValue());
+ ColumnContext columnContext = i < _numOrderByExpressions ? _orderByColumnContexts[i]
+ : resultColumnContexts.apply(_expressions.get(i));
+ columnDataTypes[i] = ColumnDataType.fromDataType(columnContext.getDataType(),
+ columnContext.isSingleValue());
}
return new DataSchema(columnNames, columnDataTypes);
}
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java
index 2950d4062bf..0c050164188 100644
--- a/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java
+++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/StreamingSelectionOrderByCombineOperatorTest.java
@@ -521,35 +521,6 @@ public class StreamingSelectionOrderByCombineOperatorTest {
"Multiset of sortedCol values must match the MinMax baseline even though the underlying rows may differ");
}
- /// Correctness under single-column ties with an OFFSET: LIMIT 10 OFFSET 20 straddles the same sortedCol=0/1
- /// boundary as [#testSingleColumnTieDeferralPreservesOrderByValuesAcrossLimitStraddle] (limit + offset = 30), but
- /// exercises it with a nonzero offset. As [#testLimitOffsetParity] documents, the server (and this combine) retains
- /// `limit + offset` rows -- the offset is trimmed by the broker afterwards -- so 30, not 10, rows come back
- /// here too; what differs from the straddle test is only the query shape, not the row count.
- @Test
- public void testSingleColumnTieDeferralPreservesOrderByValuesAcrossOffset() {
- @Language("sql") String query = "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol LIMIT 10 OFFSET 20";
- Result baseline = run(_lowCardSegments, query, false, false, false, 0);
- Result streamed = run(_lowCardSegments, query, true, false, true, 3);
- assertTrue(streamed._combineOperator instanceof StreamingSelectionOrderByCombineOperator);
- assertEquals(streamed._rows.size(), 30, "Server retains limit + offset rows; the broker trims the offset later");
- assertSorted(streamed._rows, orderByComparator(query, false));
- assertEquals(orderByColumnValues(streamed._rows), orderByColumnValues(baseline._rows),
- "Multiset of sortedCol values must match the MinMax baseline even though the underlying rows may differ");
- }
-
- /// Condition 1 (the two-expression gate): `_lowCardSegments` still tie on sortedCol alone, but a second order-by
- /// expression (valCol) makes the full order-by key a total order, so full-row parity applies unlike the
- /// single-column tests above. `_deferTiedCursors` must be false here (`orderByExpressions.size() == 1` fails), so
- /// this pins the same shape both before and after the production change: a wrongly-deferred cursor on this shape
- /// would drop a row that sorts earlier on valCol despite tying on sortedCol, which parity would catch as a missing
- /// row rather than merely a reordered one.
- @Test
- public void testTwoExpressionOrderByDoesNotDeferTiedCursors() {
- assertParity(_lowCardSegments, "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol, valCol LIMIT 30",
- false);
- }
-
/// Condition 2 (no under-delivery when the heap drains): a single-column ORDER BY with a LIMIT covering every row
/// of `_lowCardSegments` forces every segment to eventually activate no matter how aggressively ties are deferred
/// -- deferral only postpones a cursor, it never removes it from consideration, and the
@@ -1300,27 +1271,25 @@ public class StreamingSelectionOrderByCombineOperatorTest {
int blockSize = 3;
Result streamed = run(segments, query, true, nullHandling, true, blockSize);
- assertStreamingParity(streamed, baseline, comparator, query, true, blockSize);
+ assertStreamingParity(streamed, baseline, comparator, query, blockSize);
}
private void assertStreamingParity(Result result, Result baseline, Comparator<Object[]> comparator,
- @Language("sql") String query, boolean streaming, int blockSize) {
+ @Language("sql") String query, int blockSize) {
assertEquals(result._combineOperator.getClass(), StreamingSelectionOrderByCombineOperator.class,
"Expected the streaming combine operator for query: " + query);
assertEquals(result._schema, baseline._schema, "Schema mismatch for query: " + query);
assertSorted(result._rows, comparator);
assertMultisetEquals(result._rows, baseline._rows);
- if (streaming) {
- int total = 0;
- for (int size : result._blockSizes) {
- assertTrue(size > 0 && size <= blockSize,
- "Streamed block size out of range (0, " + blockSize + "] for query " + query + ": " + size);
- total += size;
- }
- assertEquals(total, result._rows.size(), "Streamed block sizes must sum to the row count for query: " + query);
- if (result._rows.size() > blockSize) {
- assertTrue(result._numBlocks >= 2, "Expected multiple streamed blocks for query: " + query);
- }
+ int total = 0;
+ for (int size : result._blockSizes) {
+ assertTrue(size > 0 && size <= blockSize,
+ "Streamed block size out of range (0, " + blockSize + "] for query " + query + ": " + size);
+ total += size;
+ }
+ assertEquals(total, result._rows.size(), "Streamed block sizes must sum to the row count for query: " + query);
+ if (result._rows.size() > blockSize) {
+ assertTrue(result._numBlocks >= 2, "Expected multiple streamed blocks for query: " + query);
}
}
@@ -1485,27 +1454,6 @@ public class StreamingSelectionOrderByCombineOperatorTest {
"The resolved mode must select the streaming combine");
}
- @Test
- public void testAutoHonoursTheMinSortedRatioThresholdForDesc() {
- // The DESC gate is a precondition, not a replacement for the ratio check: once reverse iteration is allowed, a
- // DESC query must still clear the same threshold an ASC one does. _mixedSegments is 2 sorted of 4, so 0.5 passes
- // and 0.75 does not, exactly as in testAutoHonoursTheMinSortedRatioThreshold.
- String query = "SET allowReverseOrder=true; SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol DESC, "
- + "valCol DESC LIMIT 50";
-
- QueryContext atThreshold =
- modeContext("SET sortedSelectionMergeAutoMinSortedRatio=0.5; " + query, SortedSelectionMergeMode.AUTO);
- makeStreamingInstancePlan(_mixedSegments, atThreshold);
- assertEquals(atThreshold.getSortedSelectionMergeMode(), SortedSelectionMergeMode.ON,
- "A DESC sorted ratio equal to the threshold must select the streaming merge");
-
- QueryContext aboveThreshold =
- modeContext("SET sortedSelectionMergeAutoMinSortedRatio=0.75; " + query, SortedSelectionMergeMode.AUTO);
- makeStreamingInstancePlan(_mixedSegments, aboveThreshold);
- assertEquals(aboveThreshold.getSortedSelectionMergeMode(), SortedSelectionMergeMode.OFF,
- "A DESC sorted ratio below the threshold must keep the MinMax combine");
- }
-
@Test
public void testAutoOnlyChecksTheLeadingOrderByDirection() {
// Only the leading expression rides the segment's physical order, so a DESC tail is sorted in memory and needs no
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperatorTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperatorTest.java
index 62a83bc7816..d9b0719fb4d 100644
--- a/pinot-core/src/test/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperatorTest.java
+++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/query/StreamingSelectionOrderByOperatorTest.java
@@ -20,21 +20,15 @@ package org.apache.pinot.core.operator.query;
import java.io.File;
import java.io.IOException;
-import java.util.ArrayDeque;
import java.util.ArrayList;
-import java.util.Deque;
import java.util.List;
-import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
import org.apache.commons.io.FileUtils;
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.operator.BaseProjectOperator;
-import org.apache.pinot.core.operator.ColumnContext;
-import org.apache.pinot.core.operator.ExecutionStatistics;
import org.apache.pinot.core.operator.blocks.ValueBlock;
import org.apache.pinot.core.operator.blocks.results.SelectionResultsBlock;
import org.apache.pinot.core.plan.ProjectPlanNode;
@@ -349,59 +343,20 @@ public class StreamingSelectionOrderByOperatorTest {
"Expected no nulls in the output when null handling is disabled");
}
- /// No-tail mode, single phase: [StreamingSelectionOrderByOperator] must skip a zero-document project block
- /// rather than read it as end-of-segment.
@Test
- public void testEmptyProjectBlockDoesNotTruncateSortedScan() {
- assertEmptyProjectBlocksAreSkipped(_segment, "SELECT sortedCol FROM testTable ORDER BY sortedCol LIMIT 25", 1, 2,
- 1, 0);
- }
-
- /// Same, two phase: this is also the only path that asks the injected block for `getDocIds()`.
- @Test
- public void testEmptyProjectBlockDoesNotTruncateTwoPhaseSortedScan() {
- assertEmptyProjectBlocksAreSkipped(_segment, "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol LIMIT 25",
- 1, 2, 1, 0);
- }
-
- /// Tail mode: the run scan reaches the project operator through `nextRow()`, which already tolerated empty
- /// blocks. Pins that the two scan paths agree.
- @Test
- public void testEmptyProjectBlockDoesNotTruncateRunScan() {
- assertEmptyProjectBlocksAreSkipped(_dupSegment,
- "SELECT sortedCol, tailCol FROM testTable ORDER BY sortedCol, tailCol LIMIT 25", 1, 2, 1, 0);
- }
-
- /// The skip is a loop, not a single lookahead: several empty blocks in a row must all be skipped.
- @Test
- public void testConsecutiveEmptyProjectBlocksAreAllSkipped() {
- assertEmptyProjectBlocksAreSkipped(_segment, "SELECT sortedCol FROM testTable ORDER BY sortedCol LIMIT 25", 1, 2,
- 3, 0);
- }
-
- /// Empty blocks arriving immediately before exhaustion exercise the loop's other exit: the skip must fall through
- /// to the project operator's `null` rather than spin or emit a phantom row. The limit exceeds the fixture so
- /// the scan actually reaches the end instead of stopping on the row budget.
- @Test
- public void testEmptyProjectBlocksBeforeExhaustionEndTheScan() {
- assertEmptyProjectBlocksAreSkipped(_segment, "SELECT sortedCol FROM testTable ORDER BY sortedCol LIMIT 100", 1, 0,
- 0, 2);
- }
-
- /// The skip loop sits directly on top of the `limit + offset` budget bookkeeping, so cover a non-zero offset.
- @Test
- public void testEmptyProjectBlockIsSkippedWithOffset() {
- assertEmptyProjectBlocksAreSkipped(_segment,
- "SELECT sortedCol FROM testTable ORDER BY sortedCol LIMIT 15 OFFSET 10", 1, 2, 1, 0);
- }
-
- /// `nextRow()` rebuilds the phase-1 null bitmaps in the same branch that pulls the next non-empty block, so
- /// the skip must not desynchronise a bitmap from the block it was built against. Uses the fixture that carries a
- /// null in the order-by tail column, with null handling on.
- @Test
- public void testEmptyProjectBlockIsSkippedWithNullHandling() {
- assertEmptyProjectBlocksAreSkipped(_nullTailSegment,
- "SELECT tailCol, sortedCol, valCol FROM testTable ORDER BY sortedCol, tailCol LIMIT 25", 1, true, 2, 1, 0);
+ public void testSkipsConsecutiveAndTrailingEmptyProjectBlocks() {
+ for (boolean tailToSort : new boolean[]{false, true}) {
+ IndexSegment segment = tailToSort ? _nullTailSegment : _segment;
+ String query = tailToSort
+ ? "SELECT tailCol, sortedCol, valCol FROM testTable ORDER BY sortedCol, tailCol LIMIT 100"
+ : "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol LIMIT 100";
+ List<Object[]> expected = drain(buildOperator(segment, query, 1, true, false));
+ List<Object[]> actual = drain(buildOperator(segment, query, 1, true, true));
+ assertEquals(actual.size(), expected.size());
+ for (int i = 0; i < expected.size(); i++) {
+ assertEquals(actual.get(i), expected.get(i));
+ }
+ }
}
@Test
@@ -412,39 +367,8 @@ public class StreamingSelectionOrderByOperatorTest {
@Test
public void testZeroMatchSegmentEmitsOneSchemaBlockTwoPhase() {
- assertZeroMatchEmitsOneSchemaBlock(_segment, "SELECT sortedCol, valCol FROM testTable WHERE sortedCol < 0 ORDER BY "
- + "sortedCol", "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol", 1, false);
- }
-
- @Test
- public void testZeroMatchSegmentEmitsOneSchemaBlockWithTailToSort() {
- assertZeroMatchEmitsOneSchemaBlock(_dupSegment, "SELECT sortedCol, tailCol FROM testTable WHERE sortedCol < 0 "
- + "ORDER BY sortedCol, tailCol", "SELECT sortedCol, tailCol FROM testTable ORDER BY sortedCol, tailCol", 1,
- false);
- }
-
- @Test
- public void testZeroMatchSegmentEmitsOneSchemaBlockWithNullHandling() {
- assertZeroMatchEmitsOneSchemaBlock(_nullTailSegment, "SELECT tailCol, sortedCol, valCol FROM testTable WHERE "
- + "sortedCol < 0 ORDER BY sortedCol, tailCol",
- "SELECT tailCol, sortedCol, valCol FROM testTable ORDER BY sortedCol, tailCol", 1, true);
- }
-
- @Test
- public void testMatchingSegmentEmitsNoTrailingEmptyBlock() {
- assertNoEmptyBlockAmongEmittedBlocks(_segment, "SELECT sortedCol, valCol FROM testTable ORDER BY sortedCol", 1,
- false);
- }
-
- @Test
- public void testMatchingSegmentEmitsNoTrailingEmptyBlockWithTailToSort() {
- assertNoEmptyBlockAmongEmittedBlocks(_dupSegment,
- "SELECT sortedCol, tailCol, valCol FROM testTable ORDER BY sortedCol, tailCol", 1, false);
- }
-
- @Test
- public void testMatchingSegmentEmitsNoTrailingEmptyBlockWithNullHandling() {
- assertNoEmptyBlockAmongEmittedBlocks(_nullTailSegment,
+ assertZeroMatchEmitsOneSchemaBlock(_nullTailSegment,
+ "SELECT tailCol, sortedCol, valCol FROM testTable WHERE sortedCol < 0 ORDER BY sortedCol, tailCol",
"SELECT tailCol, sortedCol, valCol FROM testTable ORDER BY sortedCol, tailCol", 1, true);
}
@@ -524,58 +448,13 @@ public class StreamingSelectionOrderByOperatorTest {
return operator;
}
- /// Every block the tail-to-sort path emits carries a structurally mutable row list. Defensive rather than a fix for
- /// a reachable failure: no consumer of this operator adds to the list today, but [SelectionResultsBlock] is a
- /// shared type and `SelectionOperatorUtils.mergeWithoutOrdering()` adds to the row list of the block it merges
- /// into, which a fixed-size `Arrays.asList` view would reject.
- @Test
- public void testRunPathEmitsGrowableRowLists() {
- QueryContext queryContext = QueryContextConverterUtils.getQueryContext(
- "SELECT sortedCol, tailCol FROM testTable ORDER BY sortedCol, tailCol");
- queryContext.setSortedSelectionMergeMode(SortedSelectionMergeMode.ON);
- Operator<SelectionResultsBlock> operator =
- new SelectionPlanNode(new SegmentContext(_dupSegment), queryContext).run();
- assertTrue(operator instanceof StreamingSelectionOrderByOperator,
- "Expected the streaming operator, got: " + operator.getClass().getSimpleName());
-
- int numBlocks = 0;
- SelectionResultsBlock block;
- while ((block = operator.nextBlock()) != null) {
- numBlocks++;
- List<Object[]> rows = block.getRows();
- assertFalse(rows.isEmpty(), "Block " + numBlocks + " must carry rows, else the assertion below is vacuous");
- int numRows = rows.size();
- rows.add(new Object[]{0, 0});
- assertEquals(rows.size(), numRows + 1, "Block " + numBlocks + " did not accept an appended row");
- }
- // One block per run, so more than one block proves the tail-to-sort path (and therefore drainAscending) ran; the
- // no-tail path emits a single block for this fixture.
- assertTrue(numBlocks > 1, "Fixture must exercise the tail-to-sort run path, got " + numBlocks + " block(s)");
- }
-
- /// The counterpart to the zero-match tests: a segment that does match rows must never emit an empty block, so the
- /// guarantee reads "at least one block" and not "always an extra one".
- private void assertNoEmptyBlockAmongEmittedBlocks(IndexSegment segment, @Language("sql") String query,
- int numSortedExpressions, boolean nullHandling) {
- StreamingSelectionOrderByOperator operator =
- buildOperatorWithInjectableProject(segment, query, numSortedExpressions, nullHandling, 0, 0, 0);
- int numBlocks = 0;
- SelectionResultsBlock block;
- while ((block = operator.nextBlock()) != null) {
- numBlocks++;
- assertFalse(block.getRows().isEmpty(),
- "Block " + numBlocks + " is empty: the schema block must only be emitted when no rows were emitted at all");
- }
- assertTrue(numBlocks > 1, "Fixture must span more than one block, otherwise a trailing block cannot be observed");
- }
-
/// A segment matching no rows must still emit exactly one block before signalling exhaustion: empty, but carrying
/// the same [DataSchema] the identical query produces when it does match rows. Consumers therefore never have
/// to reconstruct a schema the segment already knows.
private void assertZeroMatchEmitsOneSchemaBlock(IndexSegment segment, @Language("sql") String zeroMatchQuery,
@Language("sql") String matchingQuery, int numSortedExpressions, boolean nullHandling) {
SelectionResultsBlock matchingBlock =
- buildOperatorWithInjectableProject(segment, matchingQuery, numSortedExpressions, nullHandling, 0, 0, 0)
+ buildOperator(segment, matchingQuery, numSortedExpressions, nullHandling, false)
.nextBlock();
assertNotNull(matchingBlock);
assertFalse(matchingBlock.getRows().isEmpty(), "The control query must actually match rows");
@@ -583,7 +462,7 @@ public class StreamingSelectionOrderByOperatorTest {
assertNotNull(expectedSchema);
StreamingSelectionOrderByOperator operator =
- buildOperatorWithInjectableProject(segment, zeroMatchQuery, numSortedExpressions, nullHandling, 0, 0, 0);
+ buildOperator(segment, zeroMatchQuery, numSortedExpressions, nullHandling, false);
SelectionResultsBlock block = operator.nextBlock();
assertNotNull(block, "A zero-match segment must still emit one block carrying the schema");
assertTrue(block.getRows().isEmpty(), "A zero-match segment must not emit rows");
@@ -592,44 +471,8 @@ public class StreamingSelectionOrderByOperatorTest {
assertNull(operator.nextBlock(), "Exactly one block may precede exhaustion");
}
- /// Drives the operator twice over `segment`: once against the real project operator, once against one that
- /// splices zero-document blocks into its output, and asserts the emitted rows are identical.
- ///
- /// No doc-id-set operator produces an empty block today, so the case has to be injected. The project operator is
- /// built with a deliberately small `maxDocsPerCall` so the scan spans several blocks.
- ///
- /// @param injectBeforeRealBlock 1-based index of the real block to splice empties in front of; 0 for none
- /// @param numEmptyBlocks how many consecutive empty blocks to splice in at that point
- /// @param numTrailingEmptyBlocks how many empty blocks to emit after the last real block, before exhaustion
- private void assertEmptyProjectBlocksAreSkipped(IndexSegment segment, @Language("sql") String query,
- int numSortedExpressions, int injectBeforeRealBlock, int numEmptyBlocks, int numTrailingEmptyBlocks) {
- assertEmptyProjectBlocksAreSkipped(segment, query, numSortedExpressions, false, injectBeforeRealBlock,
- numEmptyBlocks, numTrailingEmptyBlocks);
- }
-
- /// As above, with explicit control over null handling.
- private void assertEmptyProjectBlocksAreSkipped(IndexSegment segment, @Language("sql") String query,
- int numSortedExpressions, boolean nullHandling, int injectBeforeRealBlock, int numEmptyBlocks,
- int numTrailingEmptyBlocks) {
- List<Object[]> expected =
- drain(buildOperatorWithInjectableProject(segment, query, numSortedExpressions, nullHandling, 0, 0, 0));
- assertTrue(expected.size() > MAX_DOCS_PER_PROJECT_BLOCK,
- "Fixture must span more than one project block, otherwise the injection is not exercised");
- List<Object[]> actual = drain(
- buildOperatorWithInjectableProject(segment, query, numSortedExpressions, nullHandling, injectBeforeRealBlock,
- numEmptyBlocks, numTrailingEmptyBlocks));
- assertEquals(actual.size(), expected.size(),
- "Row count changed when empty project blocks were injected, so the scan was truncated");
- for (int i = 0; i < expected.size(); i++) {
- assertEquals(actual.get(i), expected.get(i), "Row " + i + " mismatch for query: " + query);
- }
- }
-
- /// Builds a [StreamingSelectionOrderByOperator] directly (rather than through [SelectionPlanNode]) so
- /// the project operator can be wrapped. Passing 0 for every injection parameter yields the undecorated operator.
- private StreamingSelectionOrderByOperator buildOperatorWithInjectableProject(IndexSegment segment,
- @Language("sql") String query, int numSortedExpressions, boolean nullHandling, int injectBeforeRealBlock,
- int numEmptyBlocks, int numTrailingEmptyBlocks) {
+ private StreamingSelectionOrderByOperator buildOperator(IndexSegment segment,
+ @Language("sql") String query, int numSortedExpressions, boolean nullHandling, boolean injectEmptyBlocks) {
QueryContext queryContext = QueryContextConverterUtils.getQueryContext(query);
queryContext.setNullHandlingEnabled(nullHandling);
queryContext.setSortedSelectionMergeMode(SortedSelectionMergeMode.ON);
@@ -647,11 +490,20 @@ public class StreamingSelectionOrderByOperatorTest {
BaseProjectOperator<?> projectOperator =
new ProjectPlanNode(new SegmentContext(segment), queryContext, projectExpressions,
MAX_DOCS_PER_PROJECT_BLOCK).run();
- BaseProjectOperator<?> effective =
- injectBeforeRealBlock == 0 && numTrailingEmptyBlocks == 0
- ? projectOperator
- : new EmptyBlockInjectingProjectOperator(projectOperator, injectBeforeRealBlock, numEmptyBlocks,
- numTrailingEmptyBlocks);
+ BaseProjectOperator<?> effective = projectOperator;
+ if (injectEmptyBlocks) {
+ effective = mock(BaseProjectOperator.class, delegatesTo(projectOperator));
+ ValueBlock emptyBlock = mock(ValueBlock.class);
+ AtomicInteger leading = new AtomicInteger();
+ AtomicInteger trailing = new AtomicInteger();
+ doAnswer(invocation -> {
+ if (leading.getAndIncrement() < 2) {
+ return emptyBlock;
+ }
+ ValueBlock block = projectOperator.nextBlock();
+ return block == null && trailing.getAndIncrement() < 2 ? emptyBlock : block;
+ }).when(effective).nextBlock();
+ }
return new StreamingSelectionOrderByOperator(segment, queryContext, expressions, effective, numSortedExpressions);
}
@@ -664,133 +516,6 @@ public class StreamingSelectionOrderByOperatorTest {
return rows;
}
- /// Forwards everything to a real project operator, but splices zero-document blocks into its output:
- /// `numEmptyBlocks` of them just before the `injectBeforeRealBlock`-th real block, and
- /// `numTrailingEmptyBlocks` after the last real block but before exhaustion.
- ///
- /// Every real block is still delivered, merely deferred, so the decorator provably drops nothing of its own -
- /// any row loss observed by a test is the operator under test truncating its scan.
- private static class EmptyBlockInjectingProjectOperator extends BaseProjectOperator<ValueBlock> {
- private final BaseProjectOperator<?> _delegate;
- private final int _injectBeforeRealBlock;
- private final int _numEmptyBlocks;
- private final int _numTrailingEmptyBlocks;
- private final Deque<ValueBlock> _pending = new ArrayDeque<>();
- private int _realBlocksSeen;
- private ValueBlock _lastRealBlock;
- private boolean _injected;
- private boolean _trailingEmitted;
-
- EmptyBlockInjectingProjectOperator(BaseProjectOperator<?> delegate, int injectBeforeRealBlock, int numEmptyBlocks,
- int numTrailingEmptyBlocks) {
- _delegate = delegate;
- _injectBeforeRealBlock = injectBeforeRealBlock;
- _numEmptyBlocks = numEmptyBlocks;
- _numTrailingEmptyBlocks = numTrailingEmptyBlocks;
- }
-
- @Override
- protected ValueBlock getNextBlock() {
- if (!_pending.isEmpty()) {
- return _pending.poll();
- }
- ValueBlock block = _delegate.nextBlock();
- if (block == null) {
- // The empty blocks shadow the last real block so their value sets stay usable; there is nothing to shadow if
- // the delegate never produced one, in which case the stream simply ends.
- if (_trailingEmitted || _numTrailingEmptyBlocks == 0 || _lastRealBlock == null) {
- return null;
- }
- _trailingEmitted = true;
- for (int i = 0; i < _numTrailingEmptyBlocks; i++) {
- _pending.add(new EmptyValueBlock(_lastRealBlock));
- }
- return _pending.poll();
- }
- _lastRealBlock = block;
- _realBlocksSeen++;
- if (!_injected && _realBlocksSeen == _injectBeforeRealBlock) {
- _injected = true;
- for (int i = 0; i < _numEmptyBlocks; i++) {
- _pending.add(new EmptyValueBlock(block));
- }
- _pending.add(block);
- return _pending.poll();
- }
- return block;
- }
-
- @Override
- public Map<String, ColumnContext> getSourceColumnContextMap() {
- return _delegate.getSourceColumnContextMap();
- }
-
- @Override
- public ColumnContext getResultColumnContext(ExpressionContext expression) {
- return _delegate.getResultColumnContext(expression);
- }
-
- @Override
- public BaseProjectOperator<ValueBlock> withOrder(DocIdOrder newOrder) {
- throw new UnsupportedOperationException("Test decorator does not support reordering");
- }
-
- @Override
- public boolean isCompatibleWith(DocIdOrder order) {
- return _delegate.isCompatibleWith(order);
- }
-
- @Override
- public ExecutionStatistics getExecutionStatistics() {
- return _delegate.getExecutionStatistics();
- }
-
- @Override
- public List<? extends Operator> getChildOperators() {
- return List.of(_delegate);
- }
-
- @Override
- public String toExplainString() {
- return "EMPTY_BLOCK_INJECTING_PROJECT";
- }
- }
-
- /// A block reporting zero documents. Value sets are forwarded to the real block it shadows so that a consumer which
- /// builds its fetchers before checking the document count does not fail for the wrong reason.
- private static class EmptyValueBlock implements ValueBlock {
- private final ValueBlock _delegate;
-
- EmptyValueBlock(ValueBlock delegate) {
- _delegate = delegate;
- }
-
- @Override
- public int getNumDocs() {
- return 0;
- }
-
- @Override
- public int[] getDocIds() {
- return new int[0];
- }
-
- @Override
- public BlockValSet getBlockValueSet(ExpressionContext expression) {
- return _delegate.getBlockValueSet(expression);
- }
-
- @Override
- public BlockValSet getBlockValueSet(String column) {
- return _delegate.getBlockValueSet(column);
- }
-
- @Override
- public BlockValSet getBlockValueSet(String[] paths) {
- return _delegate.getBlockValueSet(paths);
- }
- }
-
/// Runs `query` twice over `segment` - once with the streaming hint on, once off - and asserts the
/// concatenated streaming blocks equal the materialized operator's single block, cell by cell. Returns the (now
/// exhausted) streaming operator so callers can make extra assertions on its execution statistics.
@@ -815,6 +540,7 @@ public class StreamingSelectionOrderByOperatorTest {
SelectionResultsBlock block;
while ((block = streamingOperator.nextBlock()) != null) {
numBlocks++;
+ assertFalse(block.getRows().isEmpty(), "A matching segment must not emit empty blocks");
if (streamingSchema == null) {
streamingSchema = block.getDataSchema();
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org