Inject runtime filters at the leaf-stage filter boundary

pinot-core/src/main/java/org/apache/pinot/core/plan/FilterPlanNode.java

@@ -22,6 +22,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import javax.annotation.Nullable;
 import org.apache.commons.lang3.tuple.Pair;
@@ -31,6 +32,8 @@
 import org.apache.pinot.common.request.context.predicate.JsonMatchPredicate;
 import org.apache.pinot.common.request.context.predicate.Predicate;
 import org.apache.pinot.common.request.context.predicate.RegexpLikePredicate;
+import org.apache.pinot.common.request.context.predicate.RuntimeFilter;
+import org.apache.pinot.common.request.context.predicate.RuntimeFilterPredicate;
 import org.apache.pinot.common.request.context.predicate.TextContainsPredicate;
 import org.apache.pinot.common.request.context.predicate.TextMatchPredicate;
 import org.apache.pinot.common.request.context.predicate.VectorSimilarityPredicate;
@@ -86,7 +89,43 @@ public FilterPlanNode(SegmentContext segmentContext, QueryContext queryContext,
     _indexSegment = segmentContext.getIndexSegment();
     _segmentContext = segmentContext;
     _queryContext = queryContext;
-    _filter = filter != null ? filter : _queryContext.getFilter();
+    FilterContext baseFilter = filter != null ? filter : _queryContext.getFilter();
+    _filter = combineWithRuntimeFilters(baseFilter, _queryContext);
+  }
+
+  /**
+   * If the {@link QueryContext} carries runtime filters (supplied by the MSE runtime filter
+   * feature), AND-combines a synthetic {@code RUNTIME_FILTER} predicate per column into the
+   * existing filter tree. Returns {@code baseFilter} unchanged when the map is null or empty so
+   * non-runtime-filter queries pay no extra cost.
+   *
+   * <p>The combined tree shape is:
+   * <ul>
+   *   <li>{@code baseFilter} alone, when no runtime filters are supplied.</li>
+   *   <li>A single {@code RUNTIME_FILTER} predicate (or AND of several), when the original
+   *       {@code baseFilter} is null.</li>
+   *   <li>{@code AND(baseFilter, RUNTIME_FILTER...)} when both are present.</li>
+   * </ul>
+   * The leaf evaluator validates that each runtime filter's data type matches the column's data
+   * type and throws at construction time on mismatch.</p>
+   */
+  @Nullable
+  static FilterContext combineWithRuntimeFilters(@Nullable FilterContext baseFilter, QueryContext queryContext) {
+    Map<String, RuntimeFilter> runtimeFilters = queryContext.getRuntimeFilters();
+    if (runtimeFilters == null || runtimeFilters.isEmpty()) {
+      return baseFilter;
+    }
+    List<FilterContext> injected = new ArrayList<>(runtimeFilters.size());
+    for (Map.Entry<String, RuntimeFilter> entry : runtimeFilters.entrySet()) {
+      ExpressionContext lhs = ExpressionContext.forIdentifier(entry.getKey());
+      RuntimeFilterPredicate predicate = new RuntimeFilterPredicate(lhs, entry.getValue());
+      injected.add(FilterContext.forPredicate(predicate));
+    }
+    FilterContext combinedRuntime = injected.size() == 1 ? injected.get(0) : FilterContext.forAnd(injected);
+    if (baseFilter == null) {
+      return combinedRuntime;
+    }
+    return FilterContext.forAnd(Arrays.asList(baseFilter, combinedRuntime));
   }
 
   @Override

pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java

@@ -35,6 +35,7 @@
 import org.apache.pinot.common.request.context.FunctionContext;
 import org.apache.pinot.common.request.context.OrderByExpressionContext;
 import org.apache.pinot.common.request.context.RequestContextUtils;
+import org.apache.pinot.common.request.context.predicate.RuntimeFilter;
 import org.apache.pinot.common.utils.config.QueryOptionsUtils;
 import org.apache.pinot.core.plan.maker.InstancePlanMakerImplV2;
 import org.apache.pinot.core.query.aggregation.function.AggregationFunction;
@@ -143,6 +144,10 @@ public class QueryContext {
   private boolean _accurateGroupByWithoutOrderBy;
   // Collection of index types to skip per column
   private Map<String, Set<FieldConfig.IndexType>> _skipIndexes;
+  // Runtime filters supplied by the MSE runtime filter feature, keyed by column name. When
+  // non-null and non-empty, FilterPlanNode AND-combines a synthetic RUNTIME_FILTER predicate per
+  // column into the segment-scan filter tree to prune non-matching rows during the scan.
+  private Map<String, RuntimeFilter> _runtimeFilters;
   private KeyBasedCrypterCache _crypterCache;
 
   private QueryContext(@Nullable String tableName, @Nullable QueryContext subquery,
@@ -529,6 +534,24 @@ public void setSkipIndexes(Map<String, Set<FieldConfig.IndexType>> skipIndexes)
     _skipIndexes = skipIndexes;
   }
 
+  /**
+   * Returns the runtime filters supplied by the MSE runtime filter feature, keyed by column name,
+   * or {@code null} if no runtime filters were supplied for this query.
+   */
+  @Nullable
+  public Map<String, RuntimeFilter> getRuntimeFilters() {
+    return _runtimeFilters;
+  }
+
+  /**
+   * Installs runtime filters for this query. Pass {@code null} or an empty map to leave the scan
+   * filter tree untouched. The map should be keyed by column name; each value's data type must
+   * match the column's data type or the leaf evaluator will reject it at construction time.
+   */
+  public void setRuntimeFilters(@Nullable Map<String, RuntimeFilter> runtimeFilters) {
+    _runtimeFilters = runtimeFilters;
+  }
+
   public boolean isIndexUseAllowed(String columnName, FieldConfig.IndexType indexType) {
     if (_skipIndexes == null) {
       return true;

pinot-core/src/test/java/org/apache/pinot/core/plan/FilterPlanNodeRuntimeFilterTest.java

@@ -0,0 +1,283 @@
+/**
+ * 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.plan;
+
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.pinot.common.request.context.ExpressionContext;
+import org.apache.pinot.common.request.context.FilterContext;
+import org.apache.pinot.common.request.context.predicate.BloomRuntimeFilter;
+import org.apache.pinot.common.request.context.predicate.EqPredicate;
+import org.apache.pinot.common.request.context.predicate.Predicate;
+import org.apache.pinot.common.request.context.predicate.RuntimeFilter;
+import org.apache.pinot.common.request.context.predicate.RuntimeFilterPredicate;
+import org.apache.pinot.core.query.request.context.QueryContext;
+import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+
+/**
+ * Tests {@link FilterPlanNode#combineWithRuntimeFilters} - the leaf-stage injection that
+ * AND-combines synthetic {@code RUNTIME_FILTER} predicates into the segment-scan filter tree.
+ */
+public class FilterPlanNodeRuntimeFilterTest {
+
+  private static final int EXPECTED_INSERTIONS = 1024;
+  private static final double FPP = 0.01;
+
+  /**
+   * Returns a QueryContext shaped exactly like one produced by the broker so the test exercises
+   * realistic field state (schema, expressions, etc.) instead of a hand-built skeleton.
+   */
+  private static QueryContext newQueryContext(String sql) {
+    return QueryContextConverterUtils.getQueryContext(sql);
+  }
+
+  // ---------- No-op cases (no runtime filters supplied) ----------
+
+  @Test
+  public void returnsBaseFilterWhenRuntimeMapIsNull() {
+    QueryContext qc = newQueryContext("SELECT col FROM t WHERE col = 1");
+    qc.setRuntimeFilters(null);
+    FilterContext baseFilter = qc.getFilter();
+    Assert.assertNotNull(baseFilter, "sanity: WHERE clause should produce a non-null filter");
+    Assert.assertSame(FilterPlanNode.combineWithRuntimeFilters(baseFilter, qc), baseFilter);
+  }
+
+  @Test
+  public void returnsBaseFilterWhenRuntimeMapIsEmpty() {
+    QueryContext qc = newQueryContext("SELECT col FROM t WHERE col = 1");
+    qc.setRuntimeFilters(Collections.emptyMap());
+    FilterContext baseFilter = qc.getFilter();
+    Assert.assertSame(FilterPlanNode.combineWithRuntimeFilters(baseFilter, qc), baseFilter);
+  }
+
+  @Test
+  public void returnsNullWhenBaseFilterAndRuntimeMapAreBothAbsent() {
+    QueryContext qc = newQueryContext("SELECT col FROM t");
+    qc.setRuntimeFilters(null);
+    Assert.assertNull(qc.getFilter(), "sanity: no WHERE clause means no base filter");
+    Assert.assertNull(FilterPlanNode.combineWithRuntimeFilters(null, qc));
+  }
+
+  // ---------- Injection with no original filter ----------
+
+  @Test
+  public void singleRuntimeFilterWithNullBaseProducesBarePredicateNode() {
+    QueryContext qc = newQueryContext("SELECT col FROM t");
+    qc.setRuntimeFilters(singletonRuntimeFilters("col", DataType.INT, 42));
+
+    FilterContext result = FilterPlanNode.combineWithRuntimeFilters(null, qc);
+
+    Assert.assertNotNull(result);
+    Assert.assertEquals(result.getType(), FilterContext.Type.PREDICATE,
+        "single runtime filter must not be wrapped in a redundant AND");
+    Assert.assertTrue(result.getPredicate() instanceof RuntimeFilterPredicate);
+    RuntimeFilterPredicate predicate = (RuntimeFilterPredicate) result.getPredicate();
+    Assert.assertEquals(predicate.getType(), Predicate.Type.RUNTIME_FILTER);
+    Assert.assertEquals(predicate.getLhs(), ExpressionContext.forIdentifier("col"));
+    Assert.assertEquals(predicate.getFilter().getKind(), RuntimeFilter.Kind.BLOOM);
+    Assert.assertTrue(((BloomRuntimeFilter) predicate.getFilter()).mightContain(42));
+  }
+
+  @Test
+  public void multipleRuntimeFiltersWithNullBaseAreAndCombined() {
+    Map<String, RuntimeFilter> rf = new LinkedHashMap<>();
+    rf.put("a", newBloom(DataType.INT, 1));
+    rf.put("b", newBloom(DataType.STRING, "x"));
+    QueryContext qc = newQueryContext("SELECT a, b FROM t");
+    qc.setRuntimeFilters(rf);
+
+    FilterContext result = FilterPlanNode.combineWithRuntimeFilters(null, qc);
+
+    Assert.assertNotNull(result);
+    Assert.assertEquals(result.getType(), FilterContext.Type.AND);
+    List<FilterContext> children = result.getChildren();
+    Assert.assertEquals(children.size(), 2);
+    assertRuntimeFilterPredicateOnColumn(children.get(0), "a");
+    assertRuntimeFilterPredicateOnColumn(children.get(1), "b");
+  }
+
+  // ---------- Injection with existing base filter ----------
+
+  @Test
+  public void singleRuntimeFilterIsAndedWithBaseFilter() {
+    QueryContext qc = newQueryContext("SELECT col FROM t WHERE col = 1");
+    qc.setRuntimeFilters(singletonRuntimeFilters("col", DataType.INT, 42));
+    FilterContext baseFilter = qc.getFilter();
+
+    FilterContext result = FilterPlanNode.combineWithRuntimeFilters(baseFilter, qc);
+
+    Assert.assertNotNull(result);
+    Assert.assertEquals(result.getType(), FilterContext.Type.AND);
+    List<FilterContext> children = result.getChildren();
+    Assert.assertEquals(children.size(), 2);
+    Assert.assertSame(children.get(0), baseFilter, "base filter must come first to preserve plan readability");
+    assertRuntimeFilterPredicateOnColumn(children.get(1), "col");
+  }
+
+  @Test
+  public void multipleRuntimeFiltersAreAndedWithBaseFilter() {
+    Map<String, RuntimeFilter> rf = new LinkedHashMap<>();
+    rf.put("a", newBloom(DataType.INT, 1));
+    rf.put("b", newBloom(DataType.LONG, 2L));
+    QueryContext qc = newQueryContext("SELECT a, b FROM t WHERE a > 0");
+    qc.setRuntimeFilters(rf);
+    FilterContext baseFilter = qc.getFilter();
+    Assert.assertNotNull(baseFilter);
+
+    FilterContext result = FilterPlanNode.combineWithRuntimeFilters(baseFilter, qc);
+
+    Assert.assertNotNull(result);
+    Assert.assertEquals(result.getType(), FilterContext.Type.AND);
+    Assert.assertEquals(result.getChildren().size(), 2,
+        "expected AND(baseFilter, AND(rfA, rfB)) - the runtime filters are bundled into one child");
+    Assert.assertSame(result.getChildren().get(0), baseFilter);
+    FilterContext runtimeBundle = result.getChildren().get(1);
+    Assert.assertEquals(runtimeBundle.getType(), FilterContext.Type.AND);
+    Assert.assertEquals(runtimeBundle.getChildren().size(), 2);
+  }
+
+  // ---------- Predicate properties ----------
+
+  @Test
+  public void injectedPredicateCarriesTheSuppliedFilterByReference() {
+    BloomRuntimeFilter bloom = newBloom(DataType.STRING, "user-42");
+    QueryContext qc = newQueryContext("SELECT col FROM t");
+    qc.setRuntimeFilters(Collections.singletonMap("col", bloom));
+
+    FilterContext result = FilterPlanNode.combineWithRuntimeFilters(null, qc);
+
+    Assert.assertNotNull(result);
+    RuntimeFilterPredicate predicate = (RuntimeFilterPredicate) result.getPredicate();
+    Assert.assertSame(predicate.getFilter(), bloom,
+        "the predicate must carry the same RuntimeFilter instance - reference identity matters for downstream "
+            + "evaluator wiring");
+  }
+
+  @Test
+  public void doesNotMutateTheCallersRuntimeFilterMap() {
+    Map<String, RuntimeFilter> rf = new LinkedHashMap<>();
+    rf.put("a", newBloom(DataType.INT, 1));
+    QueryContext qc = newQueryContext("SELECT a FROM t");
+    qc.setRuntimeFilters(rf);
+
+    FilterPlanNode.combineWithRuntimeFilters(null, qc);
+
+    Assert.assertEquals(rf.size(), 1);
+    Assert.assertTrue(rf.containsKey("a"));
+  }
+
+  // ---------- Coexistence with non-runtime-filter predicates in the tree ----------
+
+  @Test
+  public void existingRuntimeFilterPredicateInBaseFilterIsNotCollapsedOrDeduped() {
+    // If somebody hand-crafts a base filter that already contains a RUNTIME_FILTER predicate, the
+    // injection should still append the runtime-supplied filter rather than try to deduplicate.
+    // Deduping is the planner's job, not the leaf's.
+    QueryContext qc = newQueryContext("SELECT col FROM t");
+    BloomRuntimeFilter handCrafted = newBloom(DataType.INT, 100);
+    FilterContext handCraftedFilter = FilterContext.forPredicate(
+        new RuntimeFilterPredicate(ExpressionContext.forIdentifier("col"), handCrafted));
+
+    BloomRuntimeFilter runtimeSupplied = newBloom(DataType.INT, 200);
+    qc.setRuntimeFilters(Collections.singletonMap("col", runtimeSupplied));
+
+    FilterContext result = FilterPlanNode.combineWithRuntimeFilters(handCraftedFilter, qc);
+
+    Assert.assertNotNull(result);
+    Assert.assertEquals(result.getType(), FilterContext.Type.AND);
+    Assert.assertEquals(result.getChildren().size(), 2);
+    RuntimeFilterPredicate first = (RuntimeFilterPredicate) result.getChildren().get(0).getPredicate();
+    RuntimeFilterPredicate second = (RuntimeFilterPredicate) result.getChildren().get(1).getPredicate();
+    Assert.assertSame(first.getFilter(), handCrafted);
+    Assert.assertSame(second.getFilter(), runtimeSupplied);
+  }
+
+  // ---------- QueryContext map round-trip ----------
+
+  @Test
+  public void queryContextRoundTripsTheRuntimeFilterMap() {
+    QueryContext qc = newQueryContext("SELECT col FROM t");
+    Map<String, RuntimeFilter> rf = Collections.singletonMap("col", newBloom(DataType.INT, 7));
+    qc.setRuntimeFilters(rf);
+    Assert.assertSame(qc.getRuntimeFilters(), rf);
+  }
+
+  @Test
+  public void queryContextDefaultsRuntimeFilterMapToNull() {
+    QueryContext qc = newQueryContext("SELECT col FROM t");
+    Assert.assertNull(qc.getRuntimeFilters(),
+        "QueryContext must default to null so non-runtime-filter queries pay zero allocation");
+  }
+
+  // ---------- Sanity: ordinary predicates still work alongside ----------
+
+  @Test
+  public void baseFilterPredicateIsPreservedVerbatim() {
+    // The injection must not touch the existing predicate object - downstream code may rely on
+    // identity (e.g., for predicate-evaluator caching keyed on the Predicate instance).
+    EqPredicate eq = new EqPredicate(ExpressionContext.forIdentifier("col"), "1");
+    FilterContext baseFilter = FilterContext.forPredicate(eq);
+    QueryContext qc = newQueryContext("SELECT col FROM t");
+    qc.setRuntimeFilters(singletonRuntimeFilters("col", DataType.INT, 99));
+
+    FilterContext result = FilterPlanNode.combineWithRuntimeFilters(baseFilter, qc);
+
+    Assert.assertNotNull(result);
+    Assert.assertEquals(result.getType(), FilterContext.Type.AND);
+    Assert.assertSame(result.getChildren().get(0), baseFilter);
+    Assert.assertSame(result.getChildren().get(0).getPredicate(), eq);
+  }
+
+  // ---------- helpers ----------
+
+  private static BloomRuntimeFilter newBloom(DataType dataType, Object addedValue) {
+    BloomRuntimeFilter bf = new BloomRuntimeFilter(dataType, EXPECTED_INSERTIONS, FPP);
+    switch (dataType) {
+      case INT:
+        bf.add((int) (Integer) addedValue);
+        break;
+      case LONG:
+        bf.add((long) (Long) addedValue);
+        break;
+      case STRING:
+        bf.add((String) addedValue);
+        break;
+      default:
+        throw new IllegalArgumentException("test helper does not handle " + dataType);
+    }
+    return bf;
+  }
+
+  private static Map<String, RuntimeFilter> singletonRuntimeFilters(String column, DataType dataType,
+      Object addedValue) {
+    return Collections.singletonMap(column, newBloom(dataType, addedValue));
+  }
+
+  private static void assertRuntimeFilterPredicateOnColumn(FilterContext ctx, String expectedColumn) {
+    Assert.assertEquals(ctx.getType(), FilterContext.Type.PREDICATE);
+    Assert.assertTrue(ctx.getPredicate() instanceof RuntimeFilterPredicate);
+    Assert.assertEquals(ctx.getPredicate().getLhs(), ExpressionContext.forIdentifier(expectedColumn));
+  }
+}