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


##########
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();

Review Comment:
   Done in `c4667fa792`, though I upgraded it from an `assert` — the repo's 
surefire `argLine` has no `-ea` and I couldn't find it enabled anywhere else, 
so a Java `assert` here would never execute, in tests or in production. Since 
the failure you describe returns wrong data rather than an error, I used 
`Preconditions.checkState`. The cost is one `getCardinality()` read per query, 
against a projection pass over every surviving row. Say the word if you'd 
rather have the inert `assert` as documentation.



##########
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()));

Review Comment:
   Confirmed, and I've added a comment at that line. It isn't a new dependency: 
`ProjectPlanNode:77` resolves every selected column through this exact 
`getDataSource(column, _queryContext.getSchema())` call, as do 
`FilterPlanNode`, `DistinctPlanNode` and `AggregationPlanNode`. The late fetch 
resolves a strict subset of the columns the eager path would have resolved the 
same way, so if the schema were null the eager path would NPE first on the same 
query. The schema is only dereferenced when `getDataSourceNullable` returns 
null, i.e. for a virtual column the segment does not physically have, and it is 
set for every server query in `ServerQueryExecutorV1Impl:244`.



##########
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:
   Done in `c4667fa792`. There's now one construction loop that resolves each 
expression's context from `_columnContexts` when the scan projected it and from 
the late transform otherwise. The protected no-argument method delegates with 
an empty late expression list, so its signature is preserved and the eager path 
falls out of the same loop.



##########
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:
   Done in `c4667fa792`. Dropped `newBlockValSets()` and kept 
`getNumProjectedExpressions()`; the two subclasses now allocate the array 
directly.



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