xiangfu0 commented on code in PR #19879:
URL: https://github.com/apache/pinot/pull/19879#discussion_r4294479757


##########
pinot-core/src/main/java/org/apache/pinot/core/operator/query/LinearSelectionOrderByOperator.java:
##########
@@ -234,25 +309,122 @@ private void concatList(StringBuilder sb, List<?> list) {
 
   @Override
   protected SelectionResultsBlock getNextBlock() {
-    return new SelectionResultsBlock(createDataSchema(), 
fetch(_listBuilderSupplier), _comparator, _queryContext);
+    List<Object[]> rows = fetch(_listBuilderSupplier);
+    if (!_lateMaterialize) {
+      return new SelectionResultsBlock(createDataSchema(), rows, _comparator, 
_queryContext);
+    }
+    return fetchRemainingExpressions(rows);
   }
 
-  protected DataSchema createDataSchema() {
+  /// Fills the output expressions that were not projected during the scan, 
for the rows that survived the limit.
+  ///
+  /// The rows are visited in doc-id order because that is the order 
[BitmapDocIdSetOperator#ascending] returns, while
+  /// `rows` keeps the result order. Both views hold the same row arrays, so 
filling one fills the other.
+  private SelectionResultsBlock fetchRemainingExpressions(List<Object[]> rows) 
{
     int numExpressions = _expressions.size();
+    List<ExpressionContext> nonOrderByExpressions = 
_expressions.subList(_numOrderByExpressions, numExpressions);
+    int numNonOrderByExpressions = nonOrderByExpressions.size();
+
+    int numRows = rows.size();
+    Object[][] rowsByDocId = rows.toArray(new Object[numRows][]);
+    Arrays.sort(rowsByDocId, Comparator.comparingInt(row -> (int) 
row[_numOrderByExpressions]));
+    RoaringBitmap docIds = new RoaringBitmap();
+    for (Object[] row : rowsByDocId) {
+      docIds.add((int) row[_numOrderByExpressions]);
+    }
+    // rowsByDocId is filled by walking the bitmap in ascending order, so the 
two must line up one for one.
+    // They do only because surviving rows have distinct doc ids. If that ever 
stopped holding, the bitmap would
+    // dedupe and the trailing rows would keep their unfilled values instead 
of failing -- wrong data returned,
+    // not an error. Checked rather than asserted because this build never 
runs with -ea, and the cost is one
+    // cardinality read per query against a projection pass over every 
surviving row.
+    Preconditions.checkState(docIds.getCardinality() == numRows,

Review Comment:
   Nice -- this is exactly the guard I was hoping for, and using `checkState` 
(not `assert`) since the build runs without `-ea` is the right call. The 
per-query cardinality read is negligible.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to