Support early termination in combine operator

Author: xiangfu0

pinot-common/src/test/java/org/apache/pinot/common/utils/request/RequestUtilsTest.java

@@ -24,6 +24,7 @@
 import org.apache.calcite.sql.parser.SqlParserPos;
 import org.apache.pinot.common.request.Expression;
 import org.apache.pinot.common.request.ExpressionType;
+import org.apache.pinot.spi.utils.JsonUtils;
 import org.apache.pinot.sql.parsers.CalciteSqlParser;
 import org.apache.pinot.sql.parsers.PinotSqlType;
 import org.apache.pinot.sql.parsers.SqlNodeAndOptions;
@@ -61,6 +62,16 @@ public void testParseQuery() {
         "SELECT `foo`\n" + "FROM `countries`\n" + "WHERE `bar` > 1");
   }
 
+  @Test
+  public void testParseQueryOptionsFromJson()
+      throws Exception {
+    SqlNodeAndOptions result = RequestUtils.parseQuery("select foo from countries", JsonUtils.stringToJsonNode(
+        "{\"sql\":\"select foo from countries\","
+            + "\"queryOptions\":\"maxRowsInDistinct=5;numRowsWithoutChangeInDistinct=10\"}"));
+    assertEquals(result.getOptions().get("maxRowsInDistinct"), "5");
+    assertEquals(result.getOptions().get("numRowsWithoutChangeInDistinct"), "10");
+  }
+
   @DataProvider(name = "queryProvider")
   public Object[][] queryProvider() {
     return new Object[][] {

pinot-core/src/main/java/org/apache/pinot/core/operator/combine/DistinctCombineOperator.java

@@ -34,7 +34,7 @@ public class DistinctCombineOperator extends BaseSingleBlockCombineOperator<Dist
   private static final String EXPLAIN_NAME = "COMBINE_DISTINCT";
 
   public DistinctCombineOperator(List<Operator> operators, QueryContext queryContext, ExecutorService executorService) {
-    super(new DistinctResultsBlockMerger(), operators, queryContext, executorService);
+    super(new DistinctResultsBlockMerger(queryContext), operators, queryContext, executorService);
   }
 
   @Override

pinot-core/src/main/java/org/apache/pinot/core/operator/combine/merger/DistinctResultsBlockMerger.java

@@ -18,18 +18,53 @@
  */
 package org.apache.pinot.core.operator.combine.merger;
 
+import org.apache.pinot.common.utils.config.QueryOptionsUtils;
+import org.apache.pinot.core.operator.blocks.results.BaseResultsBlock;
 import org.apache.pinot.core.operator.blocks.results.DistinctResultsBlock;
+import org.apache.pinot.core.query.request.context.QueryContext;
 
 
 public class DistinctResultsBlockMerger implements ResultsBlockMerger<DistinctResultsBlock> {
+  private final int _maxRowsAcrossSegments;
+  private boolean _rowBudgetReached;
+
+  public DistinctResultsBlockMerger(QueryContext queryContext) {
+    if (queryContext.getQueryOptions() != null) {
+      Integer maxRows = QueryOptionsUtils.getMaxRowsInDistinct(queryContext.getQueryOptions());
+      _maxRowsAcrossSegments = maxRows != null ? maxRows : Integer.MAX_VALUE;
+    } else {
+      _maxRowsAcrossSegments = Integer.MAX_VALUE;
+    }
+  }
 
   @Override
   public boolean isQuerySatisfied(DistinctResultsBlock resultsBlock) {
+    if (_rowBudgetReached) {
+      return true;
+    }
+    if (_maxRowsAcrossSegments != Integer.MAX_VALUE
+        && resultsBlock.getDistinctTable().size() >= _maxRowsAcrossSegments) {
+      if (resultsBlock.getEarlyTerminationReason() == BaseResultsBlock.EarlyTerminationReason.NONE) {
+        resultsBlock.setEarlyTerminationReason(BaseResultsBlock.EarlyTerminationReason.DISTINCT_MAX_ROWS);
+      }
+      _rowBudgetReached = true;
+      return true;
+    }
     return resultsBlock.getDistinctTable().isSatisfied();
   }
 
   @Override
   public void mergeResultsBlocks(DistinctResultsBlock mergedBlock, DistinctResultsBlock blockToMerge) {
+    if (_rowBudgetReached) {
+      return;
+    }
     mergedBlock.getDistinctTable().mergeDistinctTable(blockToMerge.getDistinctTable());
+    if (_maxRowsAcrossSegments != Integer.MAX_VALUE
+        && mergedBlock.getDistinctTable().size() >= _maxRowsAcrossSegments) {
+      if (mergedBlock.getEarlyTerminationReason() == BaseResultsBlock.EarlyTerminationReason.NONE) {
+        mergedBlock.setEarlyTerminationReason(BaseResultsBlock.EarlyTerminationReason.DISTINCT_MAX_ROWS);
+      }
+      _rowBudgetReached = true;
+    }
   }
 }

pinot-core/src/main/java/org/apache/pinot/core/operator/query/DictionaryBasedDistinctOperator.java

@@ -63,9 +63,10 @@ public DictionaryBasedDistinctOperator(DataSource dataSource, QueryContext query
     _queryContext = queryContext;
     if (queryContext.getQueryOptions() != null) {
       Integer maxRows = QueryOptionsUtils.getMaxRowsInDistinct(queryContext.getQueryOptions());
-      _maxRowsInDistinct = maxRows != null ? maxRows : Integer.MAX_VALUE;
+      int limit = queryContext.getLimit();
+      _maxRowsInDistinct = maxRows != null ? Math.min(limit, maxRows) : limit;
     } else {
-      _maxRowsInDistinct = Integer.MAX_VALUE;
+      _maxRowsInDistinct = queryContext.getLimit();
     }
   }
 

pinot-core/src/main/java/org/apache/pinot/core/operator/query/DistinctOperator.java

@@ -90,12 +90,8 @@ protected DistinctResultsBlock getNextBlock() {
       int distinctCountBeforeBlock = enforceNoChangeLimit ? executor.getNumDistinctRowsCollected() : -1;
       boolean satisfied = executor.process(valueBlock);
       int rowsRemainingAfter = executor.getRemainingRowsToProcess();
-      int docsProcessedForLimit;
-      if (enforceRowLimit && rowsRemainingBefore != UNLIMITED_ROWS && rowsRemainingAfter != UNLIMITED_ROWS) {
-        docsProcessedForLimit = Math.max(0, rowsRemainingBefore - rowsRemainingAfter);
-      } else {
-        docsProcessedForLimit = valueBlock.getNumDocs();
-      }
+      int docsProcessedForLimit =
+          enforceRowLimit ? Math.max(0, rowsRemainingBefore - rowsRemainingAfter) : valueBlock.getNumDocs();
       _numDocsScanned += docsProcessedForLimit;
       if (enforceRowLimit && _numDocsScanned >= _maxRowsInDistinct) {
         _hitMaxRowsLimit = true;

pinot-core/src/main/java/org/apache/pinot/core/query/distinct/BaseSingleColumnDistinctExecutor.java

@@ -30,9 +30,11 @@
  * Base implementation of {@link DistinctExecutor} for single column.
  */
 public abstract class BaseSingleColumnDistinctExecutor<T extends DistinctTable, S, M> implements DistinctExecutor {
+  private static final int UNLIMITED_ROWS = Integer.MAX_VALUE;
+
   protected final ExpressionContext _expression;
   protected final T _distinctTable;
-  private int _rowsRemaining = Integer.MAX_VALUE;
+  private int _rowsRemaining = UNLIMITED_ROWS;
 
   public BaseSingleColumnDistinctExecutor(ExpressionContext expression, T distinctTable) {
     _expression = expression;
@@ -64,21 +66,26 @@ public boolean process(ValueBlock valueBlock) {
   }
 
   private boolean processWithNull(BlockValSet blockValueSet, int numDocs, RoaringBitmap nullBitmap) {
+    int limitedNumDocs = clampToRemaining(0, numDocs);
+    if (limitedNumDocs <= 0) {
+      return true;
+    }
     _distinctTable.addNull();
     S values = getValuesSV(blockValueSet);
     PeekableIntIterator nullIterator = nullBitmap.getIntIterator();
     int prev = 0;
-    while (nullIterator.hasNext()) {
+    while (nullIterator.hasNext() && prev < limitedNumDocs) {
       int nextNull = nullIterator.next();
       if (nextNull > prev) {
-        if (processSVRange(values, prev, nextNull)) {
+        int rangeEnd = Math.min(nextNull, limitedNumDocs);
+        if (processSVRange(values, prev, rangeEnd)) {
           return true;
         }
       }
       prev = nextNull + 1;
     }
-    if (prev < numDocs) {
-      return processSVRange(values, prev, numDocs);
+    if (prev < limitedNumDocs) {
+      return processSVRange(values, prev, limitedNumDocs);
     }
     return false;
   }
@@ -158,7 +165,7 @@ public int getRemainingRowsToProcess() {
   }
 
   private int clampToRemaining(int from, int to) {
-    if (_rowsRemaining == Integer.MAX_VALUE) {
+    if (_rowsRemaining == UNLIMITED_ROWS) {
       return to;
     }
     if (_rowsRemaining <= 0) {
@@ -168,7 +175,7 @@ private int clampToRemaining(int from, int to) {
   }
 
   private void consumeRows(int count) {
-    if (_rowsRemaining != Integer.MAX_VALUE) {
+    if (_rowsRemaining != UNLIMITED_ROWS) {
       _rowsRemaining -= count;
     }
   }

pinot-core/src/main/java/org/apache/pinot/core/query/distinct/DistinctExecutorFactory.java

@@ -55,7 +55,18 @@ private DistinctExecutorFactory() {
   public static DistinctExecutor getDistinctExecutor(BaseProjectOperator<?> projectOperator,
       QueryContext queryContext) {
     List<ExpressionContext> expressions = queryContext.getSelectExpressions();
-    int limit = queryContext.getLimit();
+    int queryLimit = queryContext.getLimit();
+    // The distinct table capacity is independent from the early-termination row budgets enforced by DistinctOperator.
+    // Keep a dedicated variable so it is clear we are only constraining how many unique rows we can store.
+    int tableCapacity = queryLimit;
+    if (queryContext.getQueryOptions() != null) {
+      Integer maxRowsInDistinct =
+          org.apache.pinot.common.utils.config.QueryOptionsUtils.getMaxRowsInDistinct(queryContext.getQueryOptions());
+      if (maxRowsInDistinct != null) {
+        // Only retain enough space to hold what the server will be allowed to scan.
+        tableCapacity = Math.min(queryLimit, maxRowsInDistinct);
+      }
+    }
     boolean nullHandlingEnabled = queryContext.isNullHandlingEnabled();
     List<OrderByExpressionContext> orderByExpressions = queryContext.getOrderByExpressions();
     int numExpressions = expressions.size();
@@ -76,25 +87,32 @@ public static DistinctExecutor getDistinctExecutor(BaseProjectOperator<?> projec
       // Note: Use raw value based when ordering is needed and dictionary is not sorted (consuming segments).
       if (dictionary != null && (orderByExpression == null || dictionary.isSorted())) {
         // Dictionary based
-        return new DictionaryBasedSingleColumnDistinctExecutor(expression, dictionary, dataType, limit,
+        return new DictionaryBasedSingleColumnDistinctExecutor(expression, dictionary, dataType, tableCapacity,
             nullHandlingEnabled, orderByExpression);
       } else {
         // Raw value based
         switch (dataType.getStoredType()) {
           case INT:
-            return new IntDistinctExecutor(expression, dataType, limit, nullHandlingEnabled, orderByExpression);
+            return new IntDistinctExecutor(expression, dataType, tableCapacity, nullHandlingEnabled,
+                orderByExpression);
           case LONG:
-            return new LongDistinctExecutor(expression, dataType, limit, nullHandlingEnabled, orderByExpression);
+            return new LongDistinctExecutor(expression, dataType, tableCapacity, nullHandlingEnabled,
+                orderByExpression);
           case FLOAT:
-            return new FloatDistinctExecutor(expression, dataType, limit, nullHandlingEnabled, orderByExpression);
+            return new FloatDistinctExecutor(expression, dataType, tableCapacity, nullHandlingEnabled,
+                orderByExpression);
           case DOUBLE:
-            return new DoubleDistinctExecutor(expression, dataType, limit, nullHandlingEnabled, orderByExpression);
+            return new DoubleDistinctExecutor(expression, dataType, tableCapacity, nullHandlingEnabled,
+                orderByExpression);
           case BIG_DECIMAL:
-            return new BigDecimalDistinctExecutor(expression, dataType, limit, nullHandlingEnabled, orderByExpression);
+            return new BigDecimalDistinctExecutor(expression, dataType, tableCapacity, nullHandlingEnabled,
+                orderByExpression);
           case STRING:
-            return new StringDistinctExecutor(expression, dataType, limit, nullHandlingEnabled, orderByExpression);
+            return new StringDistinctExecutor(expression, dataType, tableCapacity, nullHandlingEnabled,
+                orderByExpression);
           case BYTES:
-            return new BytesDistinctExecutor(expression, dataType, limit, nullHandlingEnabled, orderByExpression);
+            return new BytesDistinctExecutor(expression, dataType, tableCapacity, nullHandlingEnabled,
+                orderByExpression);
           default:
             throw new IllegalStateException("Unsupported data type: " + dataType);
         }
@@ -138,11 +156,11 @@ public static DistinctExecutor getDistinctExecutor(BaseProjectOperator<?> projec
       if (dictionaryBased) {
         // Dictionary based
         return new DictionaryBasedMultiColumnDistinctExecutor(expressions, hasMVExpression, dataSchema, dictionaries,
-            limit, nullHandlingEnabled, orderByExpressions);
+            tableCapacity, nullHandlingEnabled, orderByExpressions);
       } else {
         // Raw value based
-        return new RawMultiColumnDistinctExecutor(expressions, hasMVExpression, dataSchema, limit, nullHandlingEnabled,
-            orderByExpressions);
+        return new RawMultiColumnDistinctExecutor(expressions, hasMVExpression, dataSchema, tableCapacity,
+            nullHandlingEnabled, orderByExpressions);
       }
     }
   }

pinot-core/src/test/java/org/apache/pinot/core/operator/combine/DistinctResultsBlockMergerTest.java

@@ -0,0 +1,133 @@
+/**
+ * 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.IOException;
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import org.apache.pinot.common.datatable.DataTable;
+import org.apache.pinot.common.response.broker.ResultTable;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
+import org.apache.pinot.core.operator.blocks.results.BaseResultsBlock;
+import org.apache.pinot.core.operator.blocks.results.DistinctResultsBlock;
+import org.apache.pinot.core.operator.combine.merger.DistinctResultsBlockMerger;
+import org.apache.pinot.core.query.distinct.table.DistinctTable;
+import org.apache.pinot.core.query.request.context.QueryContext;
+import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertTrue;
+
+
+public class DistinctResultsBlockMergerTest {
+
+  private static final DataSchema SCHEMA =
+      new DataSchema(new String[]{"col"}, new ColumnDataType[]{ColumnDataType.INT});
+
+  @Test
+  public void shouldRespectMaxRowsAcrossSegments() {
+    QueryContext queryContext =
+        QueryContextConverterUtils.getQueryContext("SET \"maxRowsInDistinct\"=1000; SELECT DISTINCT col FROM myTable");
+    DistinctResultsBlockMerger merger = new DistinctResultsBlockMerger(queryContext);
+
+    DistinctResultsBlock merged = new DistinctResultsBlock(fakeTable(0, 800), queryContext);
+    // First block under budget
+    assertFalse(merger.isQuerySatisfied(merged));
+
+    // Merge second block that pushes us over the server-level budget (1600 rows total)
+    DistinctResultsBlock blockToMerge = new DistinctResultsBlock(fakeTable(800, 800), queryContext);
+    merger.mergeResultsBlocks(merged, blockToMerge);
+
+    assertEquals(merged.getEarlyTerminationReason(), BaseResultsBlock.EarlyTerminationReason.DISTINCT_MAX_ROWS);
+    assertTrue(merger.isQuerySatisfied(merged), "Combine should stop once server-level budget is reached");
+    assertEquals(merged.getDistinctTable().size(), 1600);
+  }
+
+  private static DistinctTable fakeTable(int startInclusive, int count) {
+    Set<Integer> values = new HashSet<>();
+    for (int i = 0; i < count; i++) {
+      values.add(startInclusive + i);
+    }
+    return new FakeDistinctTable(values);
+  }
+
+  /**
+   * Minimal {@link DistinctTable} implementation backed by a set of integers.
+   */
+  private static class FakeDistinctTable extends DistinctTable {
+    private final Set<Integer> _values;
+
+    FakeDistinctTable(Set<Integer> values) {
+      super(SCHEMA, Integer.MAX_VALUE, false);
+      _values = values;
+    }
+
+    @Override
+    public boolean hasOrderBy() {
+      return false;
+    }
+
+    @Override
+    public void mergeDistinctTable(DistinctTable distinctTable) {
+      for (Object[] row : distinctTable.getRows()) {
+        _values.add((Integer) row[0]);
+      }
+    }
+
+    @Override
+    public boolean mergeDataTable(DataTable dataTable) {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public int size() {
+      return _values.size();
+    }
+
+    @Override
+    public boolean isSatisfied() {
+      return false;
+    }
+
+    @Override
+    public List<Object[]> getRows() {
+      List<Object[]> rows = new ArrayList<>(_values.size());
+      for (Integer v : _values) {
+        rows.add(new Object[]{v});
+      }
+      return rows;
+    }
+
+    @Override
+    public DataTable toDataTable()
+        throws IOException {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public ResultTable toResultTable() {
+      throw new UnsupportedOperationException();
+    }
+  }
+}