FILTER Clauses for Aggregates

pinot-core/src/main/java/org/apache/pinot/core/operator/blocks/FilterBlock.java

@@ -30,14 +30,25 @@
  */
 public class FilterBlock implements Block {
   private final FilterBlockDocIdSet _filterBlockDocIdSet;
+  private FilterBlockDocIdSet _nonScanFilterBlockDocIdSet;
 
   public FilterBlock(FilterBlockDocIdSet filterBlockDocIdSet) {
     _filterBlockDocIdSet = filterBlockDocIdSet;
   }
 
+  /**
+   * Pre-scans the documents if needed, and returns a non-scan-based FilterBlockDocIdSet.
+   */
+  public FilterBlockDocIdSet getNonScanFilterBLockDocIdSet() {
+    if (_nonScanFilterBlockDocIdSet == null) {
+      _nonScanFilterBlockDocIdSet = _filterBlockDocIdSet.toNonScanDocIdSet();
+    }
+    return _nonScanFilterBlockDocIdSet;
+  }
+
   @Override
   public FilterBlockDocIdSet getBlockDocIdSet() {
-    return _filterBlockDocIdSet;
+    return _nonScanFilterBlockDocIdSet != null ? _nonScanFilterBlockDocIdSet : _filterBlockDocIdSet;
   }
 
   @Override

pinot-core/src/main/java/org/apache/pinot/core/operator/docidsets/BitmapDocIdSet.java

@@ -23,17 +23,19 @@
 
 
 public class BitmapDocIdSet implements FilterBlockDocIdSet {
-  private final ImmutableRoaringBitmap _docIds;
-  private final int _numDocs;
+  private final BitmapDocIdIterator _iterator;
 
   public BitmapDocIdSet(ImmutableRoaringBitmap docIds, int numDocs) {
-    _docIds = docIds;
-    _numDocs = numDocs;
+    _iterator = new BitmapDocIdIterator(docIds, numDocs);
+  }
+
+  public BitmapDocIdSet(BitmapDocIdIterator iterator) {
+    _iterator = iterator;
   }
 
   @Override
   public BitmapDocIdIterator iterator() {
-    return new BitmapDocIdIterator(_docIds, _numDocs);
+    return _iterator;
   }
 
   @Override

pinot-core/src/main/java/org/apache/pinot/core/operator/docidsets/FilterBlockDocIdSet.java

@@ -18,7 +18,16 @@
  */
 package org.apache.pinot.core.operator.docidsets;
 
+import org.apache.pinot.core.common.BlockDocIdIterator;
 import org.apache.pinot.core.common.BlockDocIdSet;
+import org.apache.pinot.core.operator.dociditerators.AndDocIdIterator;
+import org.apache.pinot.core.operator.dociditerators.BitmapDocIdIterator;
+import org.apache.pinot.core.operator.dociditerators.OrDocIdIterator;
+import org.apache.pinot.core.operator.dociditerators.RangelessBitmapDocIdIterator;
+import org.apache.pinot.core.operator.dociditerators.ScanBasedDocIdIterator;
+import org.apache.pinot.segment.spi.Constants;
+import org.roaringbitmap.RoaringBitmapWriter;
+import org.roaringbitmap.buffer.MutableRoaringBitmap;
 
 
 /**
@@ -32,4 +41,35 @@ public interface FilterBlockDocIdSet extends BlockDocIdSet {
    * filtering phase. This method should be called after the filtering is done.
    */
   long getNumEntriesScannedInFilter();
+
+  /**
+   * For scan-based FilterBlockDocIdSet, pre-scans the documents and returns a non-scan-based FilterBlockDocIdSet.
+   */
+  default FilterBlockDocIdSet toNonScanDocIdSet() {
+    BlockDocIdIterator docIdIterator = iterator();
+
+    // NOTE: AND and OR DocIdIterator might contain scan-based DocIdIterator
+    // TODO: This scan is not counted in the execution stats
+    if (docIdIterator instanceof ScanBasedDocIdIterator || docIdIterator instanceof AndDocIdIterator
+        || docIdIterator instanceof OrDocIdIterator) {
+      RoaringBitmapWriter<MutableRoaringBitmap> bitmapWriter =
+          RoaringBitmapWriter.bufferWriter().runCompress(false).get();
+      int docId;
+      while ((docId = docIdIterator.next()) != Constants.EOF) {
+        bitmapWriter.add(docId);
+      }
+      return new RangelessBitmapDocIdSet(bitmapWriter.get());
+    }
+
+    // NOTE: AND and OR DocIdSet might return BitmapBasedDocIdIterator after processing the iterators. Create a new
+    //       DocIdSet to prevent processing the iterators again
+    if (docIdIterator instanceof RangelessBitmapDocIdIterator) {
+      return new RangelessBitmapDocIdSet((RangelessBitmapDocIdIterator) docIdIterator);
+    }
+    if (docIdIterator instanceof BitmapDocIdIterator) {
+      return new BitmapDocIdSet((BitmapDocIdIterator) docIdIterator);
+    }
+
+    return this;
+  }
 }

pinot-core/src/main/java/org/apache/pinot/core/operator/docidsets/RangelessBitmapDocIdSet.java

@@ -0,0 +1,45 @@
+/**
+ * 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.docidsets;
+
+import org.apache.pinot.core.operator.dociditerators.RangelessBitmapDocIdIterator;
+import org.roaringbitmap.buffer.ImmutableRoaringBitmap;
+
+
+public class RangelessBitmapDocIdSet implements FilterBlockDocIdSet {
+  private final RangelessBitmapDocIdIterator _iterator;
+
+  public RangelessBitmapDocIdSet(ImmutableRoaringBitmap docIds) {
+    _iterator = new RangelessBitmapDocIdIterator(docIds);
+  }
+
+  public RangelessBitmapDocIdSet(RangelessBitmapDocIdIterator iterator) {
+    _iterator = iterator;
+  }
+
+  @Override
+  public RangelessBitmapDocIdIterator iterator() {
+    return _iterator;
+  }
+
+  @Override
+  public long getNumEntriesScannedInFilter() {
+    return 0L;
+  }
+}

pinot-core/src/main/java/org/apache/pinot/core/operator/filter/CombinedFilterOperator.java

@@ -0,0 +1,66 @@
+/**
+ * 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.filter;
+
+import java.util.Arrays;
+import java.util.List;
+import org.apache.pinot.core.common.Operator;
+import org.apache.pinot.core.operator.blocks.FilterBlock;
+import org.apache.pinot.core.operator.docidsets.AndDocIdSet;
+import org.apache.pinot.core.operator.docidsets.FilterBlockDocIdSet;
+
+
+/**
+ * A combined filter operator consisting of one main filter operator and one sub filter operator. The result block is
+ * the AND result of the main and sub filter.
+ */
+public class CombinedFilterOperator extends BaseFilterOperator {
+  private static final String OPERATOR_NAME = "CombinedFilterOperator";
+  private static final String EXPLAIN_NAME = "FILTER_COMBINED";
+
+  private final BaseFilterOperator _mainFilterOperator;
+  private final BaseFilterOperator _subFilterOperator;
+
+  public CombinedFilterOperator(BaseFilterOperator mainFilterOperator, BaseFilterOperator subFilterOperator) {
+    _mainFilterOperator = mainFilterOperator;
+    _subFilterOperator = subFilterOperator;
+  }
+
+  @Override
+  public String getOperatorName() {
+    return OPERATOR_NAME;
+  }
+
+  @Override
+  public List<Operator> getChildOperators() {
+    return Arrays.asList(_mainFilterOperator, _subFilterOperator);
+  }
+
+  @Override
+  public String toExplainString() {
+    return EXPLAIN_NAME;
+  }
+
+  @Override
+  protected FilterBlock getNextBlock() {
+    FilterBlockDocIdSet mainFilterDocIdSet = _mainFilterOperator.nextBlock().getNonScanFilterBLockDocIdSet();
+    FilterBlockDocIdSet subFilterDocIdSet = _subFilterOperator.nextBlock().getBlockDocIdSet();
+    return new FilterBlock(new AndDocIdSet(Arrays.asList(mainFilterDocIdSet, subFilterDocIdSet)));
+  }
+}

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

@@ -0,0 +1,116 @@
+/**
+ * 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.util.Arrays;
+import java.util.IdentityHashMap;
+import java.util.List;
+import java.util.stream.Collectors;
+import org.apache.commons.lang3.tuple.Pair;
+import org.apache.pinot.core.common.Operator;
+import org.apache.pinot.core.operator.BaseOperator;
+import org.apache.pinot.core.operator.ExecutionStatistics;
+import org.apache.pinot.core.operator.blocks.IntermediateResultsBlock;
+import org.apache.pinot.core.operator.blocks.TransformBlock;
+import org.apache.pinot.core.operator.transform.TransformOperator;
+import org.apache.pinot.core.query.aggregation.AggregationExecutor;
+import org.apache.pinot.core.query.aggregation.DefaultAggregationExecutor;
+import org.apache.pinot.core.query.aggregation.function.AggregationFunction;
+
+
+/**
+ * This operator processes a collection of filtered (and potentially non filtered) aggregations.
+ *
+ * For a query with either all aggregations being filtered or a mix of filtered and non filtered aggregations,
+ * FilteredAggregationOperator will come into execution.
+ */
+@SuppressWarnings("rawtypes")
+public class FilteredAggregationOperator extends BaseOperator<IntermediateResultsBlock> {
+  private static final String OPERATOR_NAME = "FilteredAggregationOperator";
+  private static final String EXPLAIN_NAME = "AGGREGATE_FILTERED";
+
+  private final AggregationFunction[] _aggregationFunctions;
+  private final List<Pair<AggregationFunction[], TransformOperator>> _aggFunctionsWithTransformOperator;
+  private final long _numTotalDocs;
+
+  private long _numDocsScanned;
+  private long _numEntriesScannedInFilter;
+  private long _numEntriesScannedPostFilter;
+
+  // We can potentially do away with aggregationFunctions parameter, but its cleaner to pass it in than to construct
+  // it from aggFunctionsWithTransformOperator
+  public FilteredAggregationOperator(AggregationFunction[] aggregationFunctions,
+      List<Pair<AggregationFunction[], TransformOperator>> aggFunctionsWithTransformOperator, long numTotalDocs) {
+    _aggregationFunctions = aggregationFunctions;
+    _aggFunctionsWithTransformOperator = aggFunctionsWithTransformOperator;
+    _numTotalDocs = numTotalDocs;
+  }
+
+  @Override
+  protected IntermediateResultsBlock getNextBlock() {
+    int numAggregations = _aggregationFunctions.length;
+    Object[] result = new Object[numAggregations];
+    IdentityHashMap<AggregationFunction, Integer> resultIndexMap = new IdentityHashMap<>(numAggregations);
+    for (int i = 0; i < numAggregations; i++) {
+      resultIndexMap.put(_aggregationFunctions[i], i);
+    }
+
+    for (Pair<AggregationFunction[], TransformOperator> filteredAggregation : _aggFunctionsWithTransformOperator) {
+      AggregationFunction[] aggregationFunctions = filteredAggregation.getLeft();
+      AggregationExecutor aggregationExecutor = new DefaultAggregationExecutor(aggregationFunctions);
+      TransformOperator transformOperator = filteredAggregation.getRight();
+      TransformBlock transformBlock;
+      int numDocsScanned = 0;
+      while ((transformBlock = transformOperator.nextBlock()) != null) {
+        aggregationExecutor.aggregate(transformBlock);
+        numDocsScanned += transformBlock.getNumDocs();
+      }
+      List<Object> filteredResult = aggregationExecutor.getResult();
+
+      for (int i = 0; i < aggregationFunctions.length; i++) {
+        result[resultIndexMap.get(aggregationFunctions[i])] = filteredResult.get(i);
+      }
+      _numDocsScanned += numDocsScanned;
+      _numEntriesScannedInFilter += transformOperator.getExecutionStatistics().getNumEntriesScannedInFilter();
+      _numEntriesScannedPostFilter += (long) numDocsScanned * transformOperator.getNumColumnsProjected();
+    }
+    return new IntermediateResultsBlock(_aggregationFunctions, Arrays.asList(result), false);
+  }
+
+  @Override
+  public String getOperatorName() {
+    return OPERATOR_NAME;
+  }
+
+  @Override
+  public List<Operator> getChildOperators() {
+    return _aggFunctionsWithTransformOperator.stream().map(Pair::getRight).collect(Collectors.toList());
+  }
+
+  @Override
+  public ExecutionStatistics getExecutionStatistics() {
+    return new ExecutionStatistics(_numDocsScanned, _numEntriesScannedInFilter, _numEntriesScannedPostFilter,
+        _numTotalDocs);
+  }
+
+  @Override
+  public String toExplainString() {
+    return EXPLAIN_NAME;
+  }
+}