[WIP]: Flatten per-child unions in OrFilterOperator

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

@@ -18,6 +18,8 @@
  */
 package org.apache.pinot.core.operator.filter;
 
+import java.util.Collections;
+import java.util.List;
 import org.roaringbitmap.buffer.BufferFastAggregation;
 import org.roaringbitmap.buffer.ImmutableRoaringBitmap;
 import org.roaringbitmap.buffer.MutableRoaringBitmap;
@@ -39,6 +41,23 @@ public BitmapCollection(int numDocs, boolean inverted, ImmutableRoaringBitmap...
     _bitmaps = bitmaps;
   }
 
+  /**
+   * @return whether this collection represents the complement of the union of its bitmaps (e.g. a NOT IN / != predicate)
+   */
+  public boolean isInverted() {
+    return _inverted;
+  }
+
+  /**
+   * Appends the underlying bitmaps to the given list without materializing their union. Only meaningful when
+   * {@link #isInverted()} is {@code false}: an inverted collection represents the complement of the union of these
+   * bitmaps, so the raw bitmaps alone do not represent its matching docs. Callers must check {@link #isInverted()}
+   * before relying on the appended bitmaps.
+   */
+  public void appendBitmapsTo(List<ImmutableRoaringBitmap> target) {
+    Collections.addAll(target, _bitmaps);
+  }
+
   /**
    * Inverts the bitmaps in constant time and space.
    * @return this bitmap collection inverted.

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

@@ -107,12 +107,34 @@ public boolean canOptimizeCount() {
 
   @Override
   public int getNumMatchingDocs() {
-    if (_filterOperators.size() == 2) {
+    int numChildren = _filterOperators.size();
+    // The 2-child path uses BitmapCollection.orCardinality, which handles inversion in closed form and does not
+    // materialize a union for single-bitmap children. Keep it as-is.
+    if (numChildren == 2) {
       return _filterOperators.get(0).getBitmaps().orCardinality(_filterOperators.get(1).getBitmaps());
     }
-    ImmutableRoaringBitmap[] bitmaps = new ImmutableRoaringBitmap[_filterOperators.size()];
-    for (int i = 0; i < _filterOperators.size(); i++) {
-      bitmaps[i] = _filterOperators.get(i).getBitmaps().reduce();
+    BitmapCollection[] collections = new BitmapCollection[numChildren];
+    boolean anyInverted = false;
+    for (int i = 0; i < numChildren; i++) {
+      collections[i] = _filterOperators.get(i).getBitmaps();
+      anyInverted |= collections[i].isInverted();
+    }
+    if (!anyInverted) {
+      // The union of the children equals the union of all their per-dictId bitmaps:
+      //   (a ∪ b ∪ c) ∪ (d ∪ e) = a ∪ b ∪ c ∪ d ∪ e
+      // Flatten every child's bitmaps into a single streaming union-cardinality pass, instead of reducing each
+      // child to its own union first (which allocates and discards one intermediate bitmap per child).
+      List<ImmutableRoaringBitmap> bitmaps = new ArrayList<>();
+      for (BitmapCollection collection : collections) {
+        collection.appendBitmapsTo(bitmaps);
+      }
+      return BufferFastAggregation.orCardinality(bitmaps.toArray(new ImmutableRoaringBitmap[0]));
+    }
+    // At least one child is inverted (e.g. NOT IN / !=). Its matching set is the complement of its bitmap union, so
+    // the raw bitmaps cannot be flattened. Fall back to reducing each child (which handles inversion) and combining.
+    ImmutableRoaringBitmap[] bitmaps = new ImmutableRoaringBitmap[numChildren];
+    for (int i = 0; i < numChildren; i++) {
+      bitmaps[i] = collections[i].reduce();
     }
     return BufferFastAggregation.orCardinality(bitmaps);
   }

pinot-core/src/test/java/org/apache/pinot/core/operator/filter/OrFilterOperatorTest.java

@@ -24,10 +24,16 @@
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Random;
 import java.util.TreeSet;
 import org.apache.commons.lang3.ArrayUtils;
 import org.apache.pinot.core.common.BlockDocIdIterator;
+import org.apache.pinot.core.common.BlockDocIdSet;
+import org.apache.pinot.core.common.Operator;
+import org.apache.pinot.core.operator.docidsets.EmptyDocIdSet;
 import org.apache.pinot.segment.spi.Constants;
+import org.roaringbitmap.buffer.ImmutableRoaringBitmap;
+import org.roaringbitmap.buffer.MutableRoaringBitmap;
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
@@ -181,4 +187,127 @@ public void testOrWithAllEmpty() {
     Assert.assertEquals(TestUtils.getDocIds(orFilterOperator.getTrues()), Collections.emptyList());
     Assert.assertEquals(TestUtils.getDocIds(orFilterOperator.getFalses()), Arrays.asList(0, 1, 2, 3, 4, 5, 6, 7, 8, 9));
   }
+
+  // ---------------------------------------------------------------------------------------------------------------
+  // getNumMatchingDocs() — the count-aggregation path. Children produce BitmapCollections; the operator unions them.
+  // ---------------------------------------------------------------------------------------------------------------
+
+  @Test
+  public void testGetNumMatchingDocsTwoChildren() {
+    int numDocs = 100;
+    // 2-child path (unchanged): union {1,2,3,4,10,11,20} -> 7
+    OrFilterOperator op = new OrFilterOperator(Arrays.asList(
+        new BitmapFilterOperator(numDocs, false, bitmap(1, 2, 3), bitmap(10, 11)),
+        new BitmapFilterOperator(numDocs, false, bitmap(3, 4), bitmap(11, 20))), null, numDocs, false);
+    Assert.assertEquals(op.getNumMatchingDocs(), 7);
+  }
+
+  @Test
+  public void testGetNumMatchingDocsThreeNonInvertedChildren() {
+    int numDocs = 100;
+    // flatten path: union {1,2,3,4,10,11,20,21} -> 8
+    OrFilterOperator op = new OrFilterOperator(Arrays.asList(
+        new BitmapFilterOperator(numDocs, false, bitmap(1, 2, 3), bitmap(10, 11)),
+        new BitmapFilterOperator(numDocs, false, bitmap(3, 4), bitmap(11, 20)),
+        new BitmapFilterOperator(numDocs, false, bitmap(20, 21), bitmap(2))), null, numDocs, false);
+    Assert.assertEquals(op.getNumMatchingDocs(), 8);
+  }
+
+  @Test
+  public void testGetNumMatchingDocsSingleBitmapChildren() {
+    int numDocs = 100;
+    // flatten path with one bitmap per child: union {1,2,3,4,5,6} -> 6
+    OrFilterOperator op = new OrFilterOperator(Arrays.asList(
+        new BitmapFilterOperator(numDocs, false, bitmap(1, 2, 3)),
+        new BitmapFilterOperator(numDocs, false, bitmap(3, 4, 5)),
+        new BitmapFilterOperator(numDocs, false, bitmap(5, 6))), null, numDocs, false);
+    Assert.assertEquals(op.getNumMatchingDocs(), 6);
+  }
+
+  @Test
+  public void testGetNumMatchingDocsInvertedChildFallback() {
+    int numDocs = 10; // docs 0..9
+    // c0 = {0,1}; c1 = {2}; c2 inverted over {8,9} -> matches {0..7}
+    // union = {0..7} -> 8. Exercises the inverted-child fallback (cannot flatten).
+    OrFilterOperator op = new OrFilterOperator(Arrays.asList(
+        new BitmapFilterOperator(numDocs, false, bitmap(0, 1)),
+        new BitmapFilterOperator(numDocs, false, bitmap(2)),
+        new BitmapFilterOperator(numDocs, true, bitmap(8, 9))), null, numDocs, false);
+    Assert.assertEquals(op.getNumMatchingDocs(), 8);
+  }
+
+  @Test
+  public void testGetNumMatchingDocsEquivalenceFuzz() {
+    Random random = new Random(12345);
+    int numDocs = 5000;
+    for (int trial = 0; trial < 50; trial++) {
+      int numChildren = 3 + random.nextInt(6); // 3..8 -> exercises the flatten path
+      List<BaseFilterOperator> children = new ArrayList<>(numChildren);
+      MutableRoaringBitmap expected = new MutableRoaringBitmap();
+      for (int c = 0; c < numChildren; c++) {
+        int k = 1 + random.nextInt(5);
+        ImmutableRoaringBitmap[] bitmaps = new ImmutableRoaringBitmap[k];
+        for (int j = 0; j < k; j++) {
+          MutableRoaringBitmap b = new MutableRoaringBitmap();
+          int n = 1 + random.nextInt(50);
+          for (int x = 0; x < n; x++) {
+            int doc = random.nextInt(numDocs);
+            b.add(doc);
+            expected.add(doc);
+          }
+          bitmaps[j] = b;
+        }
+        children.add(new BitmapFilterOperator(numDocs, false, bitmaps));
+      }
+      OrFilterOperator op = new OrFilterOperator(children, null, numDocs, false);
+      Assert.assertEquals(op.getNumMatchingDocs(), expected.getCardinality(),
+          "trial=" + trial + " numChildren=" + numChildren);
+    }
+  }
+
+  private static ImmutableRoaringBitmap bitmap(int... docIds) {
+    MutableRoaringBitmap bitmap = new MutableRoaringBitmap();
+    for (int docId : docIds) {
+      bitmap.add(docId);
+    }
+    return bitmap;
+  }
+
+  /** A filter operator that returns a configured {@link BitmapCollection}, to exercise the count path. */
+  private static class BitmapFilterOperator extends BaseFilterOperator {
+    private final boolean _inverted;
+    private final ImmutableRoaringBitmap[] _bitmaps;
+
+    BitmapFilterOperator(int numDocs, boolean inverted, ImmutableRoaringBitmap... bitmaps) {
+      super(numDocs, false);
+      _inverted = inverted;
+      _bitmaps = bitmaps;
+    }
+
+    @Override
+    public boolean canProduceBitmaps() {
+      return true;
+    }
+
+    @Override
+    public BitmapCollection getBitmaps() {
+      return new BitmapCollection(_numDocs, _inverted, _bitmaps);
+    }
+
+    @Override
+    protected BlockDocIdSet getTrues() {
+      // Not exercised by the getNumMatchingDocs() tests above.
+      return EmptyDocIdSet.getInstance();
+    }
+
+    @Override
+    public String toExplainString() {
+      return "FILTER_BITMAP_TEST";
+    }
+
+    @Override
+    public List<Operator> getChildOperators() {
+      return Collections.emptyList();
+    }
+  }
 }

pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkOrFilterGetNumMatchingDocs.java

@@ -0,0 +1,153 @@
+/**
+ * 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.perf;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Warmup;
+import org.openjdk.jmh.runner.Runner;
+import org.openjdk.jmh.runner.options.OptionsBuilder;
+import org.roaringbitmap.buffer.BufferFastAggregation;
+import org.roaringbitmap.buffer.ImmutableRoaringBitmap;
+import org.roaringbitmap.buffer.MutableRoaringBitmap;
+
+
+/**
+ * Isolates the work performed by {@code OrFilterOperator.getNumMatchingDocs()} on the 3+ non-inverted-children
+ * path, for a query like {@code SELECT COUNT(*) WHERE col1 IN (...) OR col2 IN (...) OR col3 IN (...)}.
+ *
+ * <p>Each child is an {@code IN} predicate that produces a {@code BitmapCollection} of K per-dictId bitmaps.
+ *
+ * <p>Two variants are compared on the same input:
+ * <ul>
+ *   <li>{@code oldPerChildReduce} — the baseline: reduce each child to its union (one materialized bitmap per
+ *       child via {@link BufferFastAggregation#or}), then {@link BufferFastAggregation#orCardinality} across the
+ *       N child unions.</li>
+ *   <li>{@code newFlatten} — the change: collect every child's bitmaps into one flat array and call
+ *       {@link BufferFastAggregation#orCardinality} exactly once. No per-child materialization.</li>
+ * </ul>
+ *
+ * <p>Both end in {@code orCardinality}; the only difference is whether each child's union is materialized first.
+ *
+ * <p>Run: {@code java -jar pinot-perf/target/benchmarks.jar BenchmarkOrFilterGetNumMatchingDocs -prof gc}
+ */
+@State(Scope.Benchmark)
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.MICROSECONDS)
+@Fork(value = 1, warmups = 0)
+@Warmup(iterations = 2, time = 2)
+@Measurement(iterations = 3, time = 2)
+public class BenchmarkOrFilterGetNumMatchingDocs {
+
+  @Param({"5000000"})
+  int _numDocs;
+
+  @Param({"1024", "100000"})
+  int _dictionaryCardinality;
+
+  /** Number of OR branches (children). */
+  @Param({"3", "5", "8"})
+  int _numChildren;
+
+  /** Width of each child's IN clause (bitmaps per child). */
+  @Param({"16", "64", "256"})
+  int _bitmapsPerChild;
+
+  // children[c] = the K bitmaps produced by child c's BitmapCollection. Null for invalid cells (K > cardinality).
+  private ImmutableRoaringBitmap[][] _children;
+
+  @Setup(Level.Trial)
+  public void setup() {
+    if (_bitmapsPerChild > _dictionaryCardinality) {
+      _children = null;
+      return;
+    }
+    _children = new ImmutableRoaringBitmap[_numChildren][];
+    for (int c = 0; c < _numChildren; c++) {
+      MutableRoaringBitmap[] bitmaps = new MutableRoaringBitmap[_bitmapsPerChild];
+      for (int j = 0; j < _bitmapsPerChild; j++) {
+        bitmaps[j] = new MutableRoaringBitmap();
+      }
+      // Independent value assignment per child (different columns); within a child the bitmaps are disjoint (SV).
+      long seed = 0x9E3779B97F4A7C15L * (c + 1);
+      for (int doc = 0; doc < _numDocs; doc++) {
+        long h = (doc ^ seed) * 0xBF58476D1CE4E5B9L;
+        int v = (int) ((h >>> 33) % _dictionaryCardinality);
+        if (v < _bitmapsPerChild) { // this child's IN clause matches dictIds 0..K-1
+          bitmaps[v].add(doc);
+        }
+      }
+      ImmutableRoaringBitmap[] immutable = new ImmutableRoaringBitmap[_bitmapsPerChild];
+      for (int j = 0; j < _bitmapsPerChild; j++) {
+        bitmaps[j].runOptimize();
+        immutable[j] = bitmaps[j];
+      }
+      _children[c] = immutable;
+    }
+  }
+
+  /** Baseline: reduce each child to its union, then orCardinality across the child unions. */
+  @Benchmark
+  public int oldPerChildReduce() {
+    if (_children == null) {
+      return 0;
+    }
+    ImmutableRoaringBitmap[] unions = new ImmutableRoaringBitmap[_children.length];
+    for (int i = 0; i < _children.length; i++) {
+      ImmutableRoaringBitmap[] child = _children[i];
+      unions[i] = child.length == 1 ? child[0] : BufferFastAggregation.or(child);
+    }
+    return BufferFastAggregation.orCardinality(unions);
+  }
+
+  /** The change: flatten all bitmaps and take the union cardinality in one pass. */
+  @Benchmark
+  public int newFlatten() {
+    if (_children == null) {
+      return 0;
+    }
+    List<ImmutableRoaringBitmap> bitmaps = new ArrayList<>();
+    for (ImmutableRoaringBitmap[] child : _children) {
+      for (ImmutableRoaringBitmap b : child) {
+        bitmaps.add(b);
+      }
+    }
+    return BufferFastAggregation.orCardinality(bitmaps.toArray(new ImmutableRoaringBitmap[0]));
+  }
+
+  public static void main(String[] args)
+      throws Exception {
+    new Runner(new OptionsBuilder()
+        .include(BenchmarkOrFilterGetNumMatchingDocs.class.getSimpleName())
+        .build())
+        .run();
+  }
+}