[8.12] ESQL: Push CIDR_MATCH to Lucene if possible (#105061) (#105177)

(cherry picked from commit b5f4c5e)

Backport of #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

@@ -54,6 +54,14 @@ public CIDRMatch(Source source, Expression ipField, List<Expression> matches) {
         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

@@ -13,6 +13,7 @@
 import org.elasticsearch.xpack.esql.evaluator.predicate.operator.comparison.Equals;
 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;
@@ -240,6 +241,8 @@ public static boolean canPushToSource(Expression exp) {
                 if (usf instanceof RegexMatch<?> || usf instanceof IsNull || usf instanceof IsNotNull) {
                     return isAttributePushable(usf.field(), usf);
                 }
+            } else if (exp instanceof CIDRMatch cidrMatch) {
+                return isAttributePushable(cidrMatch.ipField(), cidrMatch) && Expressions.foldable(cidrMatch.matches());
             }
             return false;
         }

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

@@ -8,9 +8,12 @@
 package org.elasticsearch.xpack.esql.planner;
 
 import org.elasticsearch.xpack.esql.EsqlIllegalArgumentException;
+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;
 import org.elasticsearch.xpack.ql.expression.Expression;
+import org.elasticsearch.xpack.ql.expression.Expressions;
 import org.elasticsearch.xpack.ql.expression.FieldAttribute;
 import org.elasticsearch.xpack.ql.expression.MetadataAttribute;
 import org.elasticsearch.xpack.ql.expression.function.scalar.ScalarFunction;
@@ -19,15 +22,44 @@
 import org.elasticsearch.xpack.ql.planner.ExpressionTranslator;
 import org.elasticsearch.xpack.ql.planner.ExpressionTranslators;
 import org.elasticsearch.xpack.ql.planner.QlTranslatorHandler;
+import org.elasticsearch.xpack.ql.planner.TranslatorHandler;
 import org.elasticsearch.xpack.ql.querydsl.query.Query;
+import org.elasticsearch.xpack.ql.querydsl.query.TermsQuery;
 import org.elasticsearch.xpack.ql.type.DataType;
 
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Set;
 import java.util.function.Supplier;
 
 public final class EsqlTranslatorHandler extends QlTranslatorHandler {
+
+    public static final List<ExpressionTranslator<?>> QUERY_TRANSLATORS = List.of(
+        new ExpressionTranslators.BinaryComparisons(),
+        new ExpressionTranslators.Ranges(),
+        new ExpressionTranslators.BinaryLogic(),
+        new ExpressionTranslators.IsNulls(),
+        new ExpressionTranslators.IsNotNulls(),
+        new ExpressionTranslators.Nots(),
+        new ExpressionTranslators.Likes(),
+        new ExpressionTranslators.InComparisons(),
+        new ExpressionTranslators.StringQueries(),
+        new ExpressionTranslators.Matches(),
+        new ExpressionTranslators.MultiMatches(),
+        new Scalars()
+    );
+
     @Override
     public Query asQuery(Expression e) {
-        return ExpressionTranslators.toQuery(e, this);
+        Query translation = null;
+        for (ExpressionTranslator<?> translator : QUERY_TRANSLATORS) {
+            translation = translator.translate(e, this);
+            if (translation != null) {
+                return translation;
+            }
+        }
+
+        throw new QlIllegalArgumentException("Don't know how to translate {} {}", e.nodeName(), e);
     }
 
     @Override
@@ -56,4 +88,26 @@ public Query wrapFunctionQuery(ScalarFunction sf, Expression field, Supplier<Que
         }
         throw new EsqlIllegalArgumentException("Expected a FieldAttribute or MetadataAttribute but received [" + field + "]");
     }
+
+    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;
@@ -43,24 +44,25 @@
 import org.elasticsearch.xpack.esql.session.EsqlConfiguration;
 import org.elasticsearch.xpack.esql.stats.Metrics;
 import org.elasticsearch.xpack.esql.stats.SearchStats;
-import org.elasticsearch.xpack.esql.type.EsqlDataTypes;
 import org.elasticsearch.xpack.ql.expression.Alias;
 import org.elasticsearch.xpack.ql.expression.Expression;
 import org.elasticsearch.xpack.ql.expression.Expressions;
 import org.elasticsearch.xpack.ql.expression.ReferenceAttribute;
-import org.elasticsearch.xpack.ql.expression.function.FunctionRegistry;
 import org.elasticsearch.xpack.ql.index.EsIndex;
 import org.elasticsearch.xpack.ql.index.IndexResolution;
 import org.elasticsearch.xpack.ql.tree.Source;
 import org.elasticsearch.xpack.ql.type.DataTypes;
 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.Set;
+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;
@@ -87,9 +89,8 @@ public class LocalPhysicalPlanOptimizerTests extends ESTestCase {
     private Analyzer analyzer;
     private LogicalPlanOptimizer logicalOptimizer;
     private PhysicalPlanOptimizer physicalPlanOptimizer;
+    private EsqlFunctionRegistry functionRegistry;
     private Mapper mapper;
-    private Map<String, EsField> mapping;
-    private int allFieldRowSize;
 
     private final EsqlConfiguration config;
     private final SearchStats IS_SV_STATS = new TestSearchStats() {
@@ -118,24 +119,9 @@ public LocalPhysicalPlanOptimizerTests(String name, EsqlConfiguration config) {
     @Before
     public void init() {
         parser = new EsqlParser();
-
-        mapping = loadMapping("mapping-basic.json");
-        allFieldRowSize = mapping.values()
-            .stream()
-            .mapToInt(
-                f -> (EstimatesRowSize.estimateSize(EsqlDataTypes.widenSmallNumericTypes(f.getDataType())) + f.getProperties()
-                    .values()
-                    .stream()
-                    // check one more level since the mapping contains TEXT fields with KEYWORD multi-fields
-                    .mapToInt(x -> EstimatesRowSize.estimateSize(EsqlDataTypes.widenSmallNumericTypes(x.getDataType())))
-                    .sum())
-            )
-            .sum();
-        EsIndex test = new EsIndex("test", mapping);
-        IndexResolution getIndexResult = IndexResolution.valid(test);
         logicalOptimizer = new LogicalPlanOptimizer(new LogicalOptimizerContext(EsqlTestUtils.TEST_CFG));
         physicalPlanOptimizer = new PhysicalPlanOptimizer(new PhysicalOptimizerContext(config));
-        FunctionRegistry functionRegistry = new EsqlFunctionRegistry();
+        functionRegistry = new EsqlFunctionRegistry();
         mapper = new Mapper(functionRegistry);
         var enrichResolution = new EnrichResolution(
             Set.of(
@@ -155,11 +141,15 @@ public void init() {
             ),
             Set.of("foo")
         );
+        analyzer = makeAnalyzer("mapping-basic.json", enrichResolution);
+    }
 
-        analyzer = new Analyzer(
-            new AnalyzerContext(config, functionRegistry, getIndexResult, enrichResolution),
-            new Verifier(new Metrics())
-        );
+    private Analyzer makeAnalyzer(String mappingFileName, EnrichResolution enrichResolution) {
+        var mapping = loadMapping(mappingFileName);
+        EsIndex test = new EsIndex("test", mapping);
+        IndexResolution getIndexResult = IndexResolution.valid(test);
+
+        return new Analyzer(new AnalyzerContext(config, functionRegistry, getIndexResult, enrichResolution), new Verifier(new Metrics()));
     }
 
     /**
@@ -430,6 +420,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(Set.of(), Set.of()));
+        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()));
+    }
+
     /**
      * Expects
      * LimitExec[500[INTEGER]]
@@ -493,7 +521,11 @@ private PhysicalPlan plan(String query) {
     }
 
     private PhysicalPlan plan(String query, SearchStats stats) {
-        var physical = optimizedPlan(physicalPlan(query), stats);
+        return plan(query, stats, analyzer);
+    }
+
+    private PhysicalPlan plan(String query, SearchStats stats, Analyzer analyzer) {
+        var physical = optimizedPlan(physicalPlan(query, analyzer), stats);
         return physical;
     }
 
@@ -516,7 +548,7 @@ private PhysicalPlan optimizedPlan(PhysicalPlan plan, SearchStats searchStats) {
         return l;
     }
 
-    private PhysicalPlan physicalPlan(String query) {
+    private PhysicalPlan physicalPlan(String query, Analyzer analyzer) {
         var logical = logicalOptimizer.optimize(analyzer.analyze(parser.createStatement(query)));
         // System.out.println("Logical\n" + logical);
         var physical = mapper.map(logical);
@@ -527,4 +559,13 @@ private PhysicalPlan physicalPlan(String query) {
     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);
+    }
 }