ESQL: Push CIDR_MATCH to Lucene if possible (#105061)

Author: alex-spies

docs/changelog/105061.yaml

@@ -0,0 +1,6 @@
+pr: 105061
+summary: "ESQL: Push CIDR_MATCH to Lucene if possible"
+area: ES|QL
+type: bug
+issues:
+ - 105042

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/ip/CIDRMatch.java

@@ -61,6 +61,14 @@ public CIDRMatch(
         this.matches = matches;
     }
 
+    public Expression ipField() {
+        return ipField;
+    }
+
+    public List<Expression> matches() {
+        return matches;
+    }
+
     @Override
     public ExpressionEvaluator.Factory toEvaluator(Function<Expression, ExpressionEvaluator.Factory> toEvaluator) {
         var ipEvaluatorSupplier = toEvaluator.apply(ipField);

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizer.java

@@ -14,6 +14,7 @@
 import org.elasticsearch.xpack.esql.evaluator.predicate.operator.comparison.InsensitiveBinaryComparison;
 import org.elasticsearch.xpack.esql.evaluator.predicate.operator.comparison.NotEquals;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.Count;
+import org.elasticsearch.xpack.esql.expression.function.scalar.ip.CIDRMatch;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.In;
 import org.elasticsearch.xpack.esql.optimizer.PhysicalOptimizerRules.OptimizerRule;
 import org.elasticsearch.xpack.esql.plan.physical.AggregateExec;
@@ -251,6 +252,9 @@ public static boolean canPushToSource(Expression exp, Predicate<FieldAttribute>
                 if (usf instanceof RegexMatch<?> || usf instanceof IsNull || usf instanceof IsNotNull) {
                     return isAttributePushable(usf.field(), usf, hasIdenticalDelegate);
                 }
+            } else if (exp instanceof CIDRMatch cidrMatch) {
+                return isAttributePushable(cidrMatch.ipField(), cidrMatch, hasIdenticalDelegate)
+                    && Expressions.foldable(cidrMatch.matches());
             }
             return false;
         }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/EsqlTranslatorHandler.java

@@ -11,6 +11,7 @@
 import org.elasticsearch.common.lucene.BytesRefs;
 import org.elasticsearch.xpack.esql.EsqlIllegalArgumentException;
 import org.elasticsearch.xpack.esql.evaluator.predicate.operator.comparison.InsensitiveEquals;
+import org.elasticsearch.xpack.esql.expression.function.scalar.ip.CIDRMatch;
 import org.elasticsearch.xpack.esql.querydsl.query.SingleValueQuery;
 import org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter;
 import org.elasticsearch.xpack.ql.QlIllegalArgumentException;
@@ -37,14 +38,17 @@
 import org.elasticsearch.xpack.ql.querydsl.query.MatchAll;
 import org.elasticsearch.xpack.ql.querydsl.query.Query;
 import org.elasticsearch.xpack.ql.querydsl.query.TermQuery;
+import org.elasticsearch.xpack.ql.querydsl.query.TermsQuery;
 import org.elasticsearch.xpack.ql.tree.Source;
 import org.elasticsearch.xpack.ql.type.DataType;
 import org.elasticsearch.xpack.ql.type.DataTypes;
 import org.elasticsearch.xpack.ql.util.Check;
 
 import java.math.BigDecimal;
 import java.math.BigInteger;
+import java.util.LinkedHashSet;
 import java.util.List;
+import java.util.Set;
 import java.util.function.Supplier;
 
 import static org.elasticsearch.xpack.ql.type.DataTypes.UNSIGNED_LONG;
@@ -66,7 +70,7 @@ public final class EsqlTranslatorHandler extends QlTranslatorHandler {
         new ExpressionTranslators.StringQueries(),
         new ExpressionTranslators.Matches(),
         new ExpressionTranslators.MultiMatches(),
-        new ExpressionTranslators.Scalars()
+        new Scalars()
     );
 
     @Override
@@ -245,4 +249,26 @@ private static boolean isInRange(DataType numericFieldDataType, DataType valueDa
             return minValue.compareTo(decimalValue) <= 0 && maxValue.compareTo(decimalValue) >= 0;
         }
     }
+
+    public static class Scalars extends ExpressionTranslator<ScalarFunction> {
+        @Override
+        protected Query asQuery(ScalarFunction f, TranslatorHandler handler) {
+            return doTranslate(f, handler);
+        }
+
+        public static Query doTranslate(ScalarFunction f, TranslatorHandler handler) {
+            if (f instanceof CIDRMatch cm) {
+                if (cm.ipField() instanceof FieldAttribute fa && Expressions.foldable(cm.matches())) {
+                    String targetFieldName = handler.nameOf(fa.exactAttribute());
+                    Set<Object> set = new LinkedHashSet<>(Expressions.fold(cm.matches()));
+
+                    Query query = new TermsQuery(f.source(), targetFieldName, set);
+                    // CIDR_MATCH applies only to single values.
+                    return handler.wrapFunctionQuery(f, cm.ipField(), () -> query);
+                }
+            }
+
+            return ExpressionTranslators.Scalars.doTranslate(f, handler);
+        }
+    }
 }

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizerTests.java

@@ -9,6 +9,7 @@
 
 import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
 
+import org.elasticsearch.common.network.NetworkAddress;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.core.Tuple;
 import org.elasticsearch.index.query.QueryBuilder;
@@ -54,10 +55,13 @@
 import org.elasticsearch.xpack.ql.type.EsField;
 import org.junit.Before;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;
 
 import static java.util.Arrays.asList;
+import static org.elasticsearch.common.logging.LoggerMessageFormat.format;
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.as;
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.configuration;
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.loadMapping;
@@ -415,6 +419,44 @@ public void testIsNullPushdownFilter() {
         assertThat(query.query().toString(), is(expected.toString()));
     }
 
+    /**
+     * Expects
+     * LimitExec[500[INTEGER]]
+     * \_ExchangeExec[[],false]
+     *   \_ProjectExec[[!alias_integer, boolean{f}#4, byte{f}#5, constant_keyword-foo{f}#6, date{f}#7, double{f}#8, float{f}#9,
+     *     half_float{f}#10, integer{f}#12, ip{f}#13, keyword{f}#14, long{f}#15, scaled_float{f}#11, short{f}#17, text{f}#18,
+     *     unsigned_long{f}#16, version{f}#19, wildcard{f}#20]]
+     *     \_FieldExtractExec[!alias_integer, boolean{f}#4, byte{f}#5, constant_k..][]
+     *       \_EsQueryExec[test], query[{"esql_single_value":{"field":"ip","next":{"terms":{"ip":["127.0.0.0/24"],"boost":1.0}},"source":
+     *         "cidr_match(ip, \"127.0.0.0/24\")@1:19"}}][_doc{f}#21], limit[500], sort[] estimatedRowSize[389]
+     */
+    public void testCidrMatchPushdownFilter() {
+        var allTypeMappingAnalyzer = makeAnalyzer("mapping-ip.json", new EnrichResolution());
+        final String fieldName = "ip_addr";
+
+        int cidrBlockCount = randomIntBetween(1, 10);
+        ArrayList<String> cidrBlocks = new ArrayList<>();
+        for (int i = 0; i < cidrBlockCount; i++) {
+            cidrBlocks.add(randomCidrBlock());
+        }
+        String cidrBlocksString = cidrBlocks.stream().map((s) -> "\"" + s + "\"").collect(Collectors.joining(","));
+        String cidrMatch = format(null, "cidr_match({}, {})", fieldName, cidrBlocksString);
+
+        var query = "from test | where " + cidrMatch;
+        var plan = plan(query, EsqlTestUtils.TEST_SEARCH_STATS, allTypeMappingAnalyzer);
+
+        var limit = as(plan, LimitExec.class);
+        var exchange = as(limit.child(), ExchangeExec.class);
+        var project = as(exchange.child(), ProjectExec.class);
+        var field = as(project.child(), FieldExtractExec.class);
+        var queryExec = as(field.child(), EsQueryExec.class);
+        assertThat(queryExec.limit().fold(), is(500));
+
+        var expectedInnerQuery = QueryBuilders.termsQuery(fieldName, cidrBlocks);
+        var expectedQuery = wrapWithSingleQuery(expectedInnerQuery, fieldName, new Source(1, 18, cidrMatch));
+        assertThat(queryExec.query().toString(), is(expectedQuery.toString()));
+    }
+
     private record OutOfRangeTestCase(String fieldName, String tooLow, String tooHigh) {};
 
     public void testOutOfRangeFilterPushdown() {
@@ -621,4 +663,13 @@ private PhysicalPlan physicalPlan(String query, Analyzer analyzer) {
     protected List<String> filteredWarnings() {
         return withDefaultLimitWarning(super.filteredWarnings());
     }
+
+    private String randomCidrBlock() {
+        boolean ipv4 = randomBoolean();
+
+        String address = NetworkAddress.format(randomIp(ipv4));
+        int cidrPrefixLength = ipv4 ? randomIntBetween(0, 32) : randomIntBetween(0, 128);
+
+        return format(null, "{}/{}", address, cidrPrefixLength);
+    }
 }