Jackie-Jiang commented on code in PR #19879:
URL: https://github.com/apache/pinot/pull/19879#discussion_r4288660665


##########
pinot-core/src/main/java/org/apache/pinot/core/operator/query/LinearSelectionOrderByOperator.java:
##########
@@ -234,7 +312,105 @@ 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);
+  }
+
+  /// 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]);
+    }
+
+    Set<String> columns = new HashSet<>();
+    for (ExpressionContext expression : nonOrderByExpressions) {
+      expression.getColumns(columns);
+    }
+    Map<String, DataSource> dataSourceMap = new HashMap<>();
+    for (String column : columns) {
+      dataSourceMap.put(column, _indexSegment.getDataSource(column, 
_queryContext.getSchema()));
+    }
+    int numColumns = columns.size();
+
+    BitmapDocIdSetOperator docIdSetOperator = 
BitmapDocIdSetOperator.ascending(docIds, numRows);
+    try (ProjectionOperator projectionOperator =
+        ProjectionOperatorUtils.getProjectionOperator(dataSourceMap, 
docIdSetOperator, _queryContext)) {
+      TransformOperator transformOperator =
+          new TransformOperator(_queryContext, projectionOperator, 
nonOrderByExpressions);
+
+      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[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[] row = rowsByDocId[rowBaseId + i];
+            for (int colId = 0; colId < numNonOrderByExpressions; colId++) {
+              if (nullBitmaps[colId] != null && 
nullBitmaps[colId].contains(i)) {
+                row[_numOrderByExpressions + colId] = null;
+              }
+            }
+          }
+        }
+        long entriesScanned = (long) numDocsFetched * numColumns;
+        _numEntriesScannedPostFilterLate += entriesScanned;
+        QueryScanCostContext scanCost = getScanCostContext();
+        if (scanCost != null) {
+          scanCost.addEntriesScannedPostFilter(entriesScanned);
+        }
+        rowBaseId += numDocsFetched;
+      }
+      return new SelectionResultsBlock(createDataSchema(transformOperator, 
nonOrderByExpressions), rows, _comparator,
+          _queryContext);
+    }
+  }
+
+  /// Data schema for the late-materialised case: order-by types come from the 
project operator, the rest from the
+  /// transform operator built for the late fetch.
+  private DataSchema createDataSchema(TransformOperator transformOperator,

Review Comment:
   The two createDataSchema implementations can share one construction loop. 
Resolve the context from _columnContexts for projected expressions and from the 
late transform for the remaining expressions, then build the names and types 
once. The existing protected no-argument method can delegate if we want to 
preserve its signature.



##########
pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java:
##########
@@ -1048,6 +1048,11 @@ public static class QueryOptionKey {
 
         public static final String ALLOW_REVERSE_ORDER = "allowReverseOrder";
         public static final boolean DEFAULT_ALLOW_REVERSE_ORDER = false;
+        /// Selection with an order-by prefix on a sorted column: project only 
the order-by expressions during the
+        /// scan and fetch the remaining output expressions for the rows that 
survive the limit. Disabling it
+        /// restores reading every selected column for the whole match set.
+        public static final String LATE_MATERIALIZE_SORTED_ORDER_BY = 
"lateMaterializeSortedOrderBy";

Review Comment:
   I'd remove lateMaterializeSortedOrderBy and make this the automatic behavior 
for eligible plans. The existing unsorted ORDER BY path already 
late-materializes without a query option, and another flag adds an execution 
mode to maintain. Keep the internal planner eligibility check so scans already 
bounded by LIMIT + OFFSET can still avoid an unnecessary second pass.



##########
pinot-core/src/main/java/org/apache/pinot/core/operator/query/LinearSelectionOrderByOperator.java:
##########
@@ -124,39 +159,76 @@ public IndexSegment getIndexSegment() {
   public ExecutionStatistics getExecutionStatistics() {
     long numEntriesScannedInFilter = 
_projectOperator.getExecutionStatistics().getNumEntriesScannedInFilter();
     int numDocsScanned = getNumDocsScanned();
-    long numEntriesScannedPostFilter = (long) numDocsScanned * 
_projectOperator.getNumColumnsProjected();
+    long numEntriesScannedPostFilter =
+        (long) numDocsScanned * _projectOperator.getNumColumnsProjected() + 
_numEntriesScannedPostFilterLate;
     int numTotalDocs = _indexSegment.getSegmentMetadata().getTotalDocs();
     return new ExecutionStatistics(numDocsScanned, numEntriesScannedInFilter, 
numEntriesScannedPostFilter,
         numTotalDocs);
   }
 
+  /// Number of expressions the project operator actually produces: all of 
them, or just the order-by prefix when
+  /// materialising late.
+  protected int getNumProjectedExpressions() {
+    return _columnContexts.length;
+  }
+
+  protected BlockValSet[] newBlockValSets() {

Review Comment:
   Minor simplification: getNumProjectedExpressions() and newBlockValSets() 
form a helper chain around a single array-length allocation. Keeping just the 
count accessor and allocating the BlockValSet array in the two subclasses would 
be sufficient.



-- 
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