From f16f71e2a2bd025cb2bc63cad0670deb91017818 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Iv=C3=A1n=20Cea=20Fontenla?= Date: Wed, 29 May 2024 16:23:45 +0200 Subject: [PATCH] ESQL: Add ip_prefix function (#109070) Added ESQL function to get the prefix of an IP. It works now with both IPv4 and IPv6. For users planning to use it with mixed IPs, we may need to add a function like "is_ipv4()" first. **About the skipped test:** There's currently a "bug" in the evaluators//functions that return null. Evaluators can't handle them. We'll work on support for that in another PR. It affects other functions, like `substring()`. In this function, however, it only affects in "wrong" cases (Like an invalid prefix), so it has no impact. Fixes https://github.com/elastic/elasticsearch/issues/99064 --- .gitattributes | 5 + docs/changelog/109070.yaml | 6 + .../functions/description/ip_prefix.asciidoc | 5 + .../functions/examples/ip_prefix.asciidoc | 13 ++ .../kibana/definition/ip_prefix.json | 35 ++++ .../esql/functions/kibana/docs/ip_prefix.md | 11 + .../esql/functions/layout/ip_prefix.asciidoc | 15 ++ .../functions/parameters/ip_prefix.asciidoc | 12 ++ .../esql/functions/signature/ip_prefix.svg | 1 + .../esql/functions/types/ip_prefix.asciidoc | 9 + .../xpack/esql/CsvTestUtils.java | 2 +- .../src/main/resources/ip.csv-spec | 105 ++++++++++ .../src/main/resources/meta.csv-spec | 6 +- .../function/scalar/ip/IpPrefixEvaluator.java | 183 +++++++++++++++++ .../scalar/ip/IpPrefixOnlyV4Evaluator.java | 148 ++++++++++++++ .../xpack/esql/action/EsqlCapabilities.java | 6 + .../function/EsqlFunctionRegistry.java | 2 + .../function/scalar/ip/IpPrefix.java | 191 ++++++++++++++++++ .../xpack/esql/io/stream/PlanNamedTypes.java | 2 + .../function/scalar/ip/IpPrefixTests.java | 116 +++++++++++ 20 files changed, 871 insertions(+), 2 deletions(-) create mode 100644 docs/changelog/109070.yaml create mode 100644 docs/reference/esql/functions/description/ip_prefix.asciidoc create mode 100644 docs/reference/esql/functions/examples/ip_prefix.asciidoc create mode 100644 docs/reference/esql/functions/kibana/definition/ip_prefix.json create mode 100644 docs/reference/esql/functions/kibana/docs/ip_prefix.md create mode 100644 docs/reference/esql/functions/layout/ip_prefix.asciidoc create mode 100644 docs/reference/esql/functions/parameters/ip_prefix.asciidoc create mode 100644 docs/reference/esql/functions/signature/ip_prefix.svg create mode 100644 docs/reference/esql/functions/types/ip_prefix.asciidoc create mode 100644 x-pack/plugin/esql/src/main/generated/org/elasticsearch/xpack/esql/expression/function/scalar/ip/IpPrefixEvaluator.java create mode 100644 x-pack/plugin/esql/src/main/generated/org/elasticsearch/xpack/esql/expression/function/scalar/ip/IpPrefixOnlyV4Evaluator.java create mode 100644 x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/ip/IpPrefix.java create mode 100644 x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/ip/IpPrefixTests.java diff --git a/.gitattributes b/.gitattributes index ebe1a34db347..6a8de5462ec3 100644 --- a/.gitattributes +++ b/.gitattributes @@ -8,3 +8,8 @@ x-pack/plugin/esql/src/main/antlr/*.tokens linguist-generated=true x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/*.interp linguist-generated=true x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/EsqlBaseLexer*.java linguist-generated=true x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/EsqlBaseParser*.java linguist-generated=true +x-pack/plugin/esql/src/main/generated/** linguist-generated=true + +# ESQL functions docs are autogenerated. More information at `docs/reference/esql/functions/README.md` +docs/reference/esql/functions/*/** linguist-generated=true + diff --git a/docs/changelog/109070.yaml b/docs/changelog/109070.yaml new file mode 100644 index 000000000000..8dbc0ec1c6cf --- /dev/null +++ b/docs/changelog/109070.yaml @@ -0,0 +1,6 @@ +pr: 109070 +summary: "ESQL: Add `ip_prefix` function" +area: ES|QL +type: feature +issues: + - 99064 diff --git a/docs/reference/esql/functions/description/ip_prefix.asciidoc b/docs/reference/esql/functions/description/ip_prefix.asciidoc new file mode 100644 index 000000000000..4b7a88486dea --- /dev/null +++ b/docs/reference/esql/functions/description/ip_prefix.asciidoc @@ -0,0 +1,5 @@ +// This is generated by ESQL's AbstractFunctionTestCase. Do no edit it. See ../README.md for how to regenerate it. + +*Description* + +Truncates an IP to a given prefix length. diff --git a/docs/reference/esql/functions/examples/ip_prefix.asciidoc b/docs/reference/esql/functions/examples/ip_prefix.asciidoc new file mode 100644 index 000000000000..19f0ed266afb --- /dev/null +++ b/docs/reference/esql/functions/examples/ip_prefix.asciidoc @@ -0,0 +1,13 @@ +// This is generated by ESQL's AbstractFunctionTestCase. Do no edit it. See ../README.md for how to regenerate it. + +*Example* + +[source.merge.styled,esql] +---- +include::{esql-specs}/ip.csv-spec[tag=ipPrefix] +---- +[%header.monospaced.styled,format=dsv,separator=|] +|=== +include::{esql-specs}/ip.csv-spec[tag=ipPrefix-result] +|=== + diff --git a/docs/reference/esql/functions/kibana/definition/ip_prefix.json b/docs/reference/esql/functions/kibana/definition/ip_prefix.json new file mode 100644 index 000000000000..00c3cf75a949 --- /dev/null +++ b/docs/reference/esql/functions/kibana/definition/ip_prefix.json @@ -0,0 +1,35 @@ +{ + "comment" : "This is generated by ESQL's AbstractFunctionTestCase. Do no edit it. See ../README.md for how to regenerate it.", + "type" : "eval", + "name" : "ip_prefix", + "description" : "Truncates an IP to a given prefix length.", + "signatures" : [ + { + "params" : [ + { + "name" : "ip", + "type" : "ip", + "optional" : false, + "description" : "IP address of type `ip` (both IPv4 and IPv6 are supported)." + }, + { + "name" : "prefixLengthV4", + "type" : "integer", + "optional" : false, + "description" : "Prefix length for IPv4 addresses." + }, + { + "name" : "prefixLengthV6", + "type" : "integer", + "optional" : false, + "description" : "Prefix length for IPv6 addresses." + } + ], + "variadic" : false, + "returnType" : "ip" + } + ], + "examples" : [ + "row ip4 = to_ip(\"1.2.3.4\"), ip6 = to_ip(\"fe80::cae2:65ff:fece:feb9\")\n| eval ip4_prefix = ip_prefix(ip4, 24, 0), ip6_prefix = ip_prefix(ip6, 0, 112);" + ] +} diff --git a/docs/reference/esql/functions/kibana/docs/ip_prefix.md b/docs/reference/esql/functions/kibana/docs/ip_prefix.md new file mode 100644 index 000000000000..5c0009528bb6 --- /dev/null +++ b/docs/reference/esql/functions/kibana/docs/ip_prefix.md @@ -0,0 +1,11 @@ + + +### IP_PREFIX +Truncates an IP to a given prefix length. + +``` +row ip4 = to_ip("1.2.3.4"), ip6 = to_ip("fe80::cae2:65ff:fece:feb9") +| eval ip4_prefix = ip_prefix(ip4, 24, 0), ip6_prefix = ip_prefix(ip6, 0, 112); +``` diff --git a/docs/reference/esql/functions/layout/ip_prefix.asciidoc b/docs/reference/esql/functions/layout/ip_prefix.asciidoc new file mode 100644 index 000000000000..ca51c871daf7 --- /dev/null +++ b/docs/reference/esql/functions/layout/ip_prefix.asciidoc @@ -0,0 +1,15 @@ +// This is generated by ESQL's AbstractFunctionTestCase. Do no edit it. See ../README.md for how to regenerate it. + +[discrete] +[[esql-ip_prefix]] +=== `IP_PREFIX` + +*Syntax* + +[.text-center] +image::esql/functions/signature/ip_prefix.svg[Embedded,opts=inline] + +include::../parameters/ip_prefix.asciidoc[] +include::../description/ip_prefix.asciidoc[] +include::../types/ip_prefix.asciidoc[] +include::../examples/ip_prefix.asciidoc[] diff --git a/docs/reference/esql/functions/parameters/ip_prefix.asciidoc b/docs/reference/esql/functions/parameters/ip_prefix.asciidoc new file mode 100644 index 000000000000..945601c2476e --- /dev/null +++ b/docs/reference/esql/functions/parameters/ip_prefix.asciidoc @@ -0,0 +1,12 @@ +// This is generated by ESQL's AbstractFunctionTestCase. Do no edit it. See ../README.md for how to regenerate it. + +*Parameters* + +`ip`:: +IP address of type `ip` (both IPv4 and IPv6 are supported). + +`prefixLengthV4`:: +Prefix length for IPv4 addresses. + +`prefixLengthV6`:: +Prefix length for IPv6 addresses. diff --git a/docs/reference/esql/functions/signature/ip_prefix.svg b/docs/reference/esql/functions/signature/ip_prefix.svg new file mode 100644 index 000000000000..4699c2335746 --- /dev/null +++ b/docs/reference/esql/functions/signature/ip_prefix.svg @@ -0,0 +1 @@ +IP_PREFIX(ip,prefixLengthV4,prefixLengthV6) \ No newline at end of file diff --git a/docs/reference/esql/functions/types/ip_prefix.asciidoc b/docs/reference/esql/functions/types/ip_prefix.asciidoc new file mode 100644 index 000000000000..786d99d45d32 --- /dev/null +++ b/docs/reference/esql/functions/types/ip_prefix.asciidoc @@ -0,0 +1,9 @@ +// This is generated by ESQL's AbstractFunctionTestCase. Do no edit it. See ../README.md for how to regenerate it. + +*Supported types* + +[%header.monospaced.styled,format=dsv,separator=|] +|=== +ip | prefixLengthV4 | prefixLengthV6 | result +ip | integer | integer | ip +|=== diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/CsvTestUtils.java b/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/CsvTestUtils.java index 1927cfd03ac0..7d1c168bd203 100644 --- a/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/CsvTestUtils.java +++ b/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/CsvTestUtils.java @@ -101,7 +101,7 @@ public final class CsvTestUtils { Map pairs = extractInstructions(testName); String versionRange = pairs.get("skip"); if (versionRange != null) { - String[] skipVersions = versionRange.split("-"); + String[] skipVersions = versionRange.split("-", Integer.MAX_VALUE); if (skipVersions.length != 2) { throw new IllegalArgumentException("malformed version range : " + versionRange); } diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/ip.csv-spec b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/ip.csv-spec index ae683acbb2c3..0de64d3e2d9d 100644 --- a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/ip.csv-spec +++ b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/ip.csv-spec @@ -485,3 +485,108 @@ beta | 127.0.0.1 beta | 127.0.0.1 beta | 127.0.0.1 ; + +ipPrefix +required_capability: fn_ip_prefix +//tag::ipPrefix[] +row ip4 = to_ip("1.2.3.4"), ip6 = to_ip("fe80::cae2:65ff:fece:feb9") +| eval ip4_prefix = ip_prefix(ip4, 24, 0), ip6_prefix = ip_prefix(ip6, 0, 112); +//end::ipPrefix[] + +//tag::ipPrefix-result[] +ip4:ip | ip6:ip | ip4_prefix:ip | ip6_prefix:ip +1.2.3.4 | fe80::cae2:65ff:fece:feb9 | 1.2.3.0 | fe80::cae2:65ff:fece:0000 +//end::ipPrefix-result[] +; + +ipPrefixCompleteIp +required_capability: fn_ip_prefix +row ip4 = to_ip("1.2.3.4"), ip6 = to_ip("fe80::cae2:65ff:fece:feb9") +| eval ip4_prefix = ip_prefix(ip4, 32, 0), ip6_prefix = ip_prefix(ip6, 0, 128); + +ip4:ip | ip6:ip | ip4_prefix:ip | ip6_prefix:ip +1.2.3.4 | fe80::cae2:65ff:fece:feb9 | 1.2.3.4 | fe80::cae2:65ff:fece:feb9 +; + +ipPrefixZeroBits +required_capability: fn_ip_prefix +row ip4 = to_ip("1.2.3.4"), ip6 = to_ip("fe80::cae2:65ff:fece:feb9") +| eval ip4_prefix = ip_prefix(ip4, 0, 128), ip6_prefix = ip_prefix(ip6, 32, 0); + +ip4:ip | ip6:ip | ip4_prefix:ip | ip6_prefix:ip +1.2.3.4 | fe80::cae2:65ff:fece:feb9 | 0.0.0.0 | ::0 +; + +ipPrefixWithBits +required_capability: fn_ip_prefix +row ip4 = to_ip("1.2.3.255"), ip6 = to_ip("fe80::cae2:65ff:fece:feff") +| eval ip4_prefix = ip_prefix(ip4, 25, 0), ip6_prefix = ip_prefix(ip6, 0, 121); + +ip4:ip | ip6:ip | ip4_prefix:ip | ip6_prefix:ip +1.2.3.255 | fe80::cae2:65ff:fece:feff | 1.2.3.128 | fe80::cae2:65ff:fece:fe80 +; + +ipPrefixLengthFromColumn +required_capability: fn_ip_prefix +from hosts +| where host == "alpha" +| sort card +| eval prefix = ip_prefix(ip0, 24, 128) +| keep card, ip0, prefix; + +card:keyword | ip0:ip | prefix:ip +eth0 | 127.0.0.1 | 127.0.0.0 +eth1 | ::1 | ::0 +; + +ipPrefixLengthFromExpression +required_capability: fn_ip_prefix +row ip4 = to_ip("1.2.3.4"), ip6 = to_ip("fe80::cae2:65ff:fece:feb9"), bits_per_byte = 8 +| eval ip4_length = 3 * bits_per_byte, ip4_prefix = ip_prefix(ip4, ip4_length, 0), ip6_prefix = ip_prefix(ip6, 0, 12 * 10); + +ip4:ip | ip6:ip | bits_per_byte:integer | ip4_length:integer | ip4_prefix:ip | ip6_prefix:ip +1.2.3.4 | fe80::cae2:65ff:fece:feb9 | 8 | 24 | 1.2.3.0 | fe80::cae2:65ff:fece:fe00 +; + +ipPrefixAsGroup +required_capability: fn_ip_prefix +from hosts +| stats count(*) by ip_prefix(ip1, 24, 120) +| sort `ip_prefix(ip1, 24, 120)`; +warning:Line 2:21: evaluation of [ip_prefix(ip1, 24, 120)] failed, treating result as null. Only first 20 failures recorded. +warning:Line 2:21: java.lang.IllegalArgumentException: single-value function encountered multi-value + +count(*):long | ip_prefix(ip1, 24, 120):ip +2 | ::0 +3 | 127.0.0.0 +1 | 128.0.0.0 +1 | fe80::cae2:65ff:fece:fe00 +1 | fe81::cae2:65ff:fece:fe00 +2 | null +; + +ipPrefixWithWrongLengths +required_capability: fn_ip_prefix +row ip4 = to_ip("1.2.3.4") +| eval a = ip_prefix(ip4, -1, 128), b = ip_prefix(ip4, 32, -1), c = ip_prefix(ip4, 33, 0), d = ip_prefix(ip4, 32, 129); +warning:Line 2:12: evaluation of [ip_prefix(ip4, -1, 128)] failed, treating result as null. Only first 20 failures recorded. +warning:Line 2:12: java.lang.IllegalArgumentException: Prefix length v4 must be in range [0, 32], found -1 +warning:Line 2:41: evaluation of [ip_prefix(ip4, 32, -1)] failed, treating result as null. Only first 20 failures recorded. +warning:Line 2:41: java.lang.IllegalArgumentException: Prefix length v6 must be in range [0, 128], found -1 +warning:Line 2:69: evaluation of [ip_prefix(ip4, 33, 0)] failed, treating result as null. Only first 20 failures recorded. +warning:Line 2:69: java.lang.IllegalArgumentException: Prefix length v4 must be in range [0, 32], found 33 +warning:Line 2:96: evaluation of [ip_prefix(ip4, 32, 129)] failed, treating result as null. Only first 20 failures recorded. +warning:Line 2:96: java.lang.IllegalArgumentException: Prefix length v6 must be in range [0, 128], found 129 + +ip4:ip | a:ip | b:ip | c:ip | d:ip +1.2.3.4 | null | null | null | null +; + +ipPrefixWithNullArguments +required_capability: fn_ip_prefix +row ip4 = to_ip("1.2.3.4") +| eval a = ip_prefix(null, 32, 128), b = ip_prefix(ip4, null, 128), c = ip_prefix(ip4, 32, null); + +ip4:ip | a:ip | b:ip | c:ip +1.2.3.4 | null | null | null +; diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/meta.csv-spec b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/meta.csv-spec index f68dd15e9c51..eff4cb05bd8c 100644 --- a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/meta.csv-spec +++ b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/meta.csv-spec @@ -30,6 +30,7 @@ double e() "double|integer|long|unsigned_long floor(number:double|integer|long|unsigned_long)" "keyword from_base64(string:keyword|text)" "boolean|double|integer|ip|keyword|long|text|version greatest(first:boolean|double|integer|ip|keyword|long|text|version, ?rest...:boolean|double|integer|ip|keyword|long|text|version)" +"ip ip_prefix(ip:ip, prefixLengthV4:integer, prefixLengthV6:integer)" "boolean|double|integer|ip|keyword|long|text|version least(first:boolean|double|integer|ip|keyword|long|text|version, ?rest...:boolean|double|integer|ip|keyword|long|text|version)" "keyword left(string:keyword|text, length:integer)" "integer length(string:keyword|text)" @@ -144,6 +145,7 @@ ends_with |[str, suffix] |["keyword|text", "keyword|te floor |number |"double|integer|long|unsigned_long" |Numeric expression. If `null`, the function returns `null`. from_base64 |string |"keyword|text" |A base64 string. greatest |first |"boolean|double|integer|ip|keyword|long|text|version" |First of the columns to evaluate. +ip_prefix |[ip, prefixLengthV4, prefixLengthV6]|[ip, integer, integer] |[IP address of type `ip` (both IPv4 and IPv6 are supported)., Prefix length for IPv4 addresses., Prefix length for IPv6 addresses.] least |first |"boolean|double|integer|ip|keyword|long|text|version" |First of the columns to evaluate. left |[string, length] |["keyword|text", integer] |[The string from which to return a substring., The number of characters to return.] length |string |"keyword|text" |String expression. If `null`, the function returns `null`. @@ -259,6 +261,7 @@ ends_with |Returns a boolean that indicates whether a keyword string ends wi floor |Round a number down to the nearest integer. from_base64 |Decode a base64 string. greatest |Returns the maximum value from multiple columns. This is similar to <> except it is intended to run on multiple columns at once. +ip_prefix |Truncates an IP to a given prefix length. least |Returns the minimum value from multiple columns. This is similar to <> except it is intended to run on multiple columns at once. left |Returns the substring that extracts 'length' chars from 'string' starting from the left. length |Returns the character length of a string. @@ -375,6 +378,7 @@ ends_with |boolean floor |"double|integer|long|unsigned_long" |false |false |false from_base64 |keyword |false |false |false greatest |"boolean|double|integer|ip|keyword|long|text|version" |false |true |false +ip_prefix |ip |[false, false, false] |false |false least |"boolean|double|integer|ip|keyword|long|text|version" |false |true |false left |keyword |[false, false] |false |false length |integer |false |false |false @@ -471,5 +475,5 @@ countFunctions#[skip:-8.14.99, reason:BIN added] meta functions | stats a = count(*), b = count(*), c = count(*) | mv_expand c; a:long | b:long | c:long -106 | 106 | 106 +107 | 107 | 107 ; diff --git a/x-pack/plugin/esql/src/main/generated/org/elasticsearch/xpack/esql/expression/function/scalar/ip/IpPrefixEvaluator.java b/x-pack/plugin/esql/src/main/generated/org/elasticsearch/xpack/esql/expression/function/scalar/ip/IpPrefixEvaluator.java new file mode 100644 index 000000000000..174df48d5ce6 --- /dev/null +++ b/x-pack/plugin/esql/src/main/generated/org/elasticsearch/xpack/esql/expression/function/scalar/ip/IpPrefixEvaluator.java @@ -0,0 +1,183 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License +// 2.0; you may not use this file except in compliance with the Elastic License +// 2.0. +package org.elasticsearch.xpack.esql.expression.function.scalar.ip; + +import java.lang.IllegalArgumentException; +import java.lang.Override; +import java.lang.String; +import java.util.function.Function; +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.data.BytesRefBlock; +import org.elasticsearch.compute.data.BytesRefVector; +import org.elasticsearch.compute.data.IntBlock; +import org.elasticsearch.compute.data.IntVector; +import org.elasticsearch.compute.data.Page; +import org.elasticsearch.compute.operator.DriverContext; +import org.elasticsearch.compute.operator.EvalOperator; +import org.elasticsearch.core.Releasables; +import org.elasticsearch.xpack.esql.core.tree.Source; +import org.elasticsearch.xpack.esql.expression.function.Warnings; + +/** + * {@link EvalOperator.ExpressionEvaluator} implementation for {@link IpPrefix}. + * This class is generated. Do not edit it. + */ +public final class IpPrefixEvaluator implements EvalOperator.ExpressionEvaluator { + private final Warnings warnings; + + private final EvalOperator.ExpressionEvaluator ip; + + private final EvalOperator.ExpressionEvaluator prefixLengthV4; + + private final EvalOperator.ExpressionEvaluator prefixLengthV6; + + private final BytesRef scratch; + + private final DriverContext driverContext; + + public IpPrefixEvaluator(Source source, EvalOperator.ExpressionEvaluator ip, + EvalOperator.ExpressionEvaluator prefixLengthV4, + EvalOperator.ExpressionEvaluator prefixLengthV6, BytesRef scratch, + DriverContext driverContext) { + this.warnings = new Warnings(source); + this.ip = ip; + this.prefixLengthV4 = prefixLengthV4; + this.prefixLengthV6 = prefixLengthV6; + this.scratch = scratch; + this.driverContext = driverContext; + } + + @Override + public Block eval(Page page) { + try (BytesRefBlock ipBlock = (BytesRefBlock) ip.eval(page)) { + try (IntBlock prefixLengthV4Block = (IntBlock) prefixLengthV4.eval(page)) { + try (IntBlock prefixLengthV6Block = (IntBlock) prefixLengthV6.eval(page)) { + BytesRefVector ipVector = ipBlock.asVector(); + if (ipVector == null) { + return eval(page.getPositionCount(), ipBlock, prefixLengthV4Block, prefixLengthV6Block); + } + IntVector prefixLengthV4Vector = prefixLengthV4Block.asVector(); + if (prefixLengthV4Vector == null) { + return eval(page.getPositionCount(), ipBlock, prefixLengthV4Block, prefixLengthV6Block); + } + IntVector prefixLengthV6Vector = prefixLengthV6Block.asVector(); + if (prefixLengthV6Vector == null) { + return eval(page.getPositionCount(), ipBlock, prefixLengthV4Block, prefixLengthV6Block); + } + return eval(page.getPositionCount(), ipVector, prefixLengthV4Vector, prefixLengthV6Vector); + } + } + } + } + + public BytesRefBlock eval(int positionCount, BytesRefBlock ipBlock, IntBlock prefixLengthV4Block, + IntBlock prefixLengthV6Block) { + try(BytesRefBlock.Builder result = driverContext.blockFactory().newBytesRefBlockBuilder(positionCount)) { + BytesRef ipScratch = new BytesRef(); + position: for (int p = 0; p < positionCount; p++) { + if (ipBlock.isNull(p)) { + result.appendNull(); + continue position; + } + if (ipBlock.getValueCount(p) != 1) { + if (ipBlock.getValueCount(p) > 1) { + warnings.registerException(new IllegalArgumentException("single-value function encountered multi-value")); + } + result.appendNull(); + continue position; + } + if (prefixLengthV4Block.isNull(p)) { + result.appendNull(); + continue position; + } + if (prefixLengthV4Block.getValueCount(p) != 1) { + if (prefixLengthV4Block.getValueCount(p) > 1) { + warnings.registerException(new IllegalArgumentException("single-value function encountered multi-value")); + } + result.appendNull(); + continue position; + } + if (prefixLengthV6Block.isNull(p)) { + result.appendNull(); + continue position; + } + if (prefixLengthV6Block.getValueCount(p) != 1) { + if (prefixLengthV6Block.getValueCount(p) > 1) { + warnings.registerException(new IllegalArgumentException("single-value function encountered multi-value")); + } + result.appendNull(); + continue position; + } + try { + result.appendBytesRef(IpPrefix.process(ipBlock.getBytesRef(ipBlock.getFirstValueIndex(p), ipScratch), prefixLengthV4Block.getInt(prefixLengthV4Block.getFirstValueIndex(p)), prefixLengthV6Block.getInt(prefixLengthV6Block.getFirstValueIndex(p)), scratch)); + } catch (IllegalArgumentException e) { + warnings.registerException(e); + result.appendNull(); + } + } + return result.build(); + } + } + + public BytesRefBlock eval(int positionCount, BytesRefVector ipVector, + IntVector prefixLengthV4Vector, IntVector prefixLengthV6Vector) { + try(BytesRefBlock.Builder result = driverContext.blockFactory().newBytesRefBlockBuilder(positionCount)) { + BytesRef ipScratch = new BytesRef(); + position: for (int p = 0; p < positionCount; p++) { + try { + result.appendBytesRef(IpPrefix.process(ipVector.getBytesRef(p, ipScratch), prefixLengthV4Vector.getInt(p), prefixLengthV6Vector.getInt(p), scratch)); + } catch (IllegalArgumentException e) { + warnings.registerException(e); + result.appendNull(); + } + } + return result.build(); + } + } + + @Override + public String toString() { + return "IpPrefixEvaluator[" + "ip=" + ip + ", prefixLengthV4=" + prefixLengthV4 + ", prefixLengthV6=" + prefixLengthV6 + "]"; + } + + @Override + public void close() { + Releasables.closeExpectNoException(ip, prefixLengthV4, prefixLengthV6); + } + + static class Factory implements EvalOperator.ExpressionEvaluator.Factory { + private final Source source; + + private final EvalOperator.ExpressionEvaluator.Factory ip; + + private final EvalOperator.ExpressionEvaluator.Factory prefixLengthV4; + + private final EvalOperator.ExpressionEvaluator.Factory prefixLengthV6; + + private final Function scratch; + + public Factory(Source source, EvalOperator.ExpressionEvaluator.Factory ip, + EvalOperator.ExpressionEvaluator.Factory prefixLengthV4, + EvalOperator.ExpressionEvaluator.Factory prefixLengthV6, + Function scratch) { + this.source = source; + this.ip = ip; + this.prefixLengthV4 = prefixLengthV4; + this.prefixLengthV6 = prefixLengthV6; + this.scratch = scratch; + } + + @Override + public IpPrefixEvaluator get(DriverContext context) { + return new IpPrefixEvaluator(source, ip.get(context), prefixLengthV4.get(context), prefixLengthV6.get(context), scratch.apply(context), context); + } + + @Override + public String toString() { + return "IpPrefixEvaluator[" + "ip=" + ip + ", prefixLengthV4=" + prefixLengthV4 + ", prefixLengthV6=" + prefixLengthV6 + "]"; + } + } +} diff --git a/x-pack/plugin/esql/src/main/generated/org/elasticsearch/xpack/esql/expression/function/scalar/ip/IpPrefixOnlyV4Evaluator.java b/x-pack/plugin/esql/src/main/generated/org/elasticsearch/xpack/esql/expression/function/scalar/ip/IpPrefixOnlyV4Evaluator.java new file mode 100644 index 000000000000..a6cb7c7f9b68 --- /dev/null +++ b/x-pack/plugin/esql/src/main/generated/org/elasticsearch/xpack/esql/expression/function/scalar/ip/IpPrefixOnlyV4Evaluator.java @@ -0,0 +1,148 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License +// 2.0; you may not use this file except in compliance with the Elastic License +// 2.0. +package org.elasticsearch.xpack.esql.expression.function.scalar.ip; + +import java.lang.IllegalArgumentException; +import java.lang.Override; +import java.lang.String; +import java.util.function.Function; +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.data.BytesRefBlock; +import org.elasticsearch.compute.data.BytesRefVector; +import org.elasticsearch.compute.data.IntBlock; +import org.elasticsearch.compute.data.IntVector; +import org.elasticsearch.compute.data.Page; +import org.elasticsearch.compute.operator.DriverContext; +import org.elasticsearch.compute.operator.EvalOperator; +import org.elasticsearch.core.Releasables; +import org.elasticsearch.xpack.esql.core.tree.Source; +import org.elasticsearch.xpack.esql.expression.function.Warnings; + +/** + * {@link EvalOperator.ExpressionEvaluator} implementation for {@link IpPrefix}. + * This class is generated. Do not edit it. + */ +public final class IpPrefixOnlyV4Evaluator implements EvalOperator.ExpressionEvaluator { + private final Warnings warnings; + + private final EvalOperator.ExpressionEvaluator ip; + + private final EvalOperator.ExpressionEvaluator prefixLengthV4; + + private final BytesRef scratch; + + private final DriverContext driverContext; + + public IpPrefixOnlyV4Evaluator(Source source, EvalOperator.ExpressionEvaluator ip, + EvalOperator.ExpressionEvaluator prefixLengthV4, BytesRef scratch, + DriverContext driverContext) { + this.warnings = new Warnings(source); + this.ip = ip; + this.prefixLengthV4 = prefixLengthV4; + this.scratch = scratch; + this.driverContext = driverContext; + } + + @Override + public Block eval(Page page) { + try (BytesRefBlock ipBlock = (BytesRefBlock) ip.eval(page)) { + try (IntBlock prefixLengthV4Block = (IntBlock) prefixLengthV4.eval(page)) { + BytesRefVector ipVector = ipBlock.asVector(); + if (ipVector == null) { + return eval(page.getPositionCount(), ipBlock, prefixLengthV4Block); + } + IntVector prefixLengthV4Vector = prefixLengthV4Block.asVector(); + if (prefixLengthV4Vector == null) { + return eval(page.getPositionCount(), ipBlock, prefixLengthV4Block); + } + return eval(page.getPositionCount(), ipVector, prefixLengthV4Vector).asBlock(); + } + } + } + + public BytesRefBlock eval(int positionCount, BytesRefBlock ipBlock, + IntBlock prefixLengthV4Block) { + try(BytesRefBlock.Builder result = driverContext.blockFactory().newBytesRefBlockBuilder(positionCount)) { + BytesRef ipScratch = new BytesRef(); + position: for (int p = 0; p < positionCount; p++) { + if (ipBlock.isNull(p)) { + result.appendNull(); + continue position; + } + if (ipBlock.getValueCount(p) != 1) { + if (ipBlock.getValueCount(p) > 1) { + warnings.registerException(new IllegalArgumentException("single-value function encountered multi-value")); + } + result.appendNull(); + continue position; + } + if (prefixLengthV4Block.isNull(p)) { + result.appendNull(); + continue position; + } + if (prefixLengthV4Block.getValueCount(p) != 1) { + if (prefixLengthV4Block.getValueCount(p) > 1) { + warnings.registerException(new IllegalArgumentException("single-value function encountered multi-value")); + } + result.appendNull(); + continue position; + } + result.appendBytesRef(IpPrefix.process(ipBlock.getBytesRef(ipBlock.getFirstValueIndex(p), ipScratch), prefixLengthV4Block.getInt(prefixLengthV4Block.getFirstValueIndex(p)), scratch)); + } + return result.build(); + } + } + + public BytesRefVector eval(int positionCount, BytesRefVector ipVector, + IntVector prefixLengthV4Vector) { + try(BytesRefVector.Builder result = driverContext.blockFactory().newBytesRefVectorBuilder(positionCount)) { + BytesRef ipScratch = new BytesRef(); + position: for (int p = 0; p < positionCount; p++) { + result.appendBytesRef(IpPrefix.process(ipVector.getBytesRef(p, ipScratch), prefixLengthV4Vector.getInt(p), scratch)); + } + return result.build(); + } + } + + @Override + public String toString() { + return "IpPrefixOnlyV4Evaluator[" + "ip=" + ip + ", prefixLengthV4=" + prefixLengthV4 + "]"; + } + + @Override + public void close() { + Releasables.closeExpectNoException(ip, prefixLengthV4); + } + + static class Factory implements EvalOperator.ExpressionEvaluator.Factory { + private final Source source; + + private final EvalOperator.ExpressionEvaluator.Factory ip; + + private final EvalOperator.ExpressionEvaluator.Factory prefixLengthV4; + + private final Function scratch; + + public Factory(Source source, EvalOperator.ExpressionEvaluator.Factory ip, + EvalOperator.ExpressionEvaluator.Factory prefixLengthV4, + Function scratch) { + this.source = source; + this.ip = ip; + this.prefixLengthV4 = prefixLengthV4; + this.scratch = scratch; + } + + @Override + public IpPrefixOnlyV4Evaluator get(DriverContext context) { + return new IpPrefixOnlyV4Evaluator(source, ip.get(context), prefixLengthV4.get(context), scratch.apply(context), context); + } + + @Override + public String toString() { + return "IpPrefixOnlyV4Evaluator[" + "ip=" + ip + ", prefixLengthV4=" + prefixLengthV4 + "]"; + } + } +} diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java index ca90b9e1e3f2..e8f136c297ce 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java @@ -27,6 +27,11 @@ public class EsqlCapabilities { */ private static final String FN_CBRT = "fn_cbrt"; + /** + * Support for function {@code IP_PREFIX}. + */ + private static final String FN_IP_PREFIX = "fn_ip_prefix"; + /** * Optimization for ST_CENTROID changed some results in cartesian data. #108713 */ @@ -47,6 +52,7 @@ public class EsqlCapabilities { private static Set capabilities() { List caps = new ArrayList<>(); caps.add(FN_CBRT); + caps.add(FN_IP_PREFIX); caps.add(ST_CENTROID_AGG_OPTIMIZED); caps.add(METADATA_IGNORED_FIELD); diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/EsqlFunctionRegistry.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/EsqlFunctionRegistry.java index 90f28b263f07..382b02d2f9b3 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/EsqlFunctionRegistry.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/EsqlFunctionRegistry.java @@ -52,6 +52,7 @@ import org.elasticsearch.xpack.esql.expression.function.scalar.date.DateParse; import org.elasticsearch.xpack.esql.expression.function.scalar.date.DateTrunc; import org.elasticsearch.xpack.esql.expression.function.scalar.date.Now; import org.elasticsearch.xpack.esql.expression.function.scalar.ip.CIDRMatch; +import org.elasticsearch.xpack.esql.expression.function.scalar.ip.IpPrefix; import org.elasticsearch.xpack.esql.expression.function.scalar.math.Abs; import org.elasticsearch.xpack.esql.expression.function.scalar.math.Acos; import org.elasticsearch.xpack.esql.expression.function.scalar.math.Asin; @@ -257,6 +258,7 @@ public final class EsqlFunctionRegistry extends FunctionRegistry { new FunctionDefinition[] { def(Coalesce.class, Coalesce::new, "coalesce"), }, // IP new FunctionDefinition[] { def(CIDRMatch.class, CIDRMatch::new, "cidr_match") }, + new FunctionDefinition[] { def(IpPrefix.class, IpPrefix::new, "ip_prefix") }, // conversion functions new FunctionDefinition[] { def(FromBase64.class, FromBase64::new, "from_base64"), diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/ip/IpPrefix.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/ip/IpPrefix.java new file mode 100644 index 000000000000..d271742aef82 --- /dev/null +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/ip/IpPrefix.java @@ -0,0 +1,191 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.esql.expression.function.scalar.ip; + +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.compute.ann.Evaluator; +import org.elasticsearch.compute.ann.Fixed; +import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator; +import org.elasticsearch.xpack.esql.core.expression.Expression; +import org.elasticsearch.xpack.esql.core.expression.Expressions; +import org.elasticsearch.xpack.esql.core.expression.function.OptionalArgument; +import org.elasticsearch.xpack.esql.core.tree.NodeInfo; +import org.elasticsearch.xpack.esql.core.tree.Source; +import org.elasticsearch.xpack.esql.core.type.DataType; +import org.elasticsearch.xpack.esql.core.type.DataTypes; +import org.elasticsearch.xpack.esql.expression.function.Example; +import org.elasticsearch.xpack.esql.expression.function.FunctionInfo; +import org.elasticsearch.xpack.esql.expression.function.Param; +import org.elasticsearch.xpack.esql.expression.function.scalar.EsqlScalarFunction; +import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput; +import org.elasticsearch.xpack.esql.io.stream.PlanStreamOutput; + +import java.io.IOException; +import java.util.Arrays; +import java.util.List; +import java.util.function.Function; + +import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.FIRST; +import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.SECOND; +import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.THIRD; +import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isIPAndExact; +import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isType; +import static org.elasticsearch.xpack.esql.core.type.DataTypes.INTEGER; + +/** + * Truncates an IP value to a given prefix length. + */ +public class IpPrefix extends EsqlScalarFunction implements OptionalArgument { + // Borrowed from Lucene, rfc4291 prefix + private static final byte[] IPV4_PREFIX = new byte[] { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -1, -1 }; + + private final Expression ipField; + private final Expression prefixLengthV4Field; + private final Expression prefixLengthV6Field; + + @FunctionInfo( + returnType = "ip", + description = "Truncates an IP to a given prefix length.", + examples = @Example(file = "ip", tag = "ipPrefix") + ) + public IpPrefix( + Source source, + @Param( + name = "ip", + type = { "ip" }, + description = "IP address of type `ip` (both IPv4 and IPv6 are supported)." + ) Expression ipField, + @Param( + name = "prefixLengthV4", + type = { "integer" }, + description = "Prefix length for IPv4 addresses." + ) Expression prefixLengthV4Field, + @Param( + name = "prefixLengthV6", + type = { "integer" }, + description = "Prefix length for IPv6 addresses." + ) Expression prefixLengthV6Field + ) { + super(source, Arrays.asList(ipField, prefixLengthV4Field, prefixLengthV6Field)); + this.ipField = ipField; + this.prefixLengthV4Field = prefixLengthV4Field; + this.prefixLengthV6Field = prefixLengthV6Field; + } + + public static IpPrefix readFrom(PlanStreamInput in) throws IOException { + return new IpPrefix(in.readSource(), in.readExpression(), in.readExpression(), in.readExpression()); + } + + public static void writeTo(PlanStreamOutput out, IpPrefix ipPrefix) throws IOException { + out.writeSource(ipPrefix.source()); + List fields = ipPrefix.children(); + assert fields.size() == 3; + out.writeExpression(fields.get(0)); + out.writeExpression(fields.get(1)); + out.writeExpression(fields.get(2)); + } + + public Expression ipField() { + return ipField; + } + + public Expression prefixLengthV4Field() { + return prefixLengthV4Field; + } + + public Expression prefixLengthV6Field() { + return prefixLengthV6Field; + } + + @Override + public boolean foldable() { + return Expressions.foldable(children()); + } + + @Override + public ExpressionEvaluator.Factory toEvaluator(Function toEvaluator) { + var ipEvaluatorSupplier = toEvaluator.apply(ipField); + var prefixLengthV4EvaluatorSupplier = toEvaluator.apply(prefixLengthV4Field); + var prefixLengthV6EvaluatorSupplier = toEvaluator.apply(prefixLengthV6Field); + + return new IpPrefixEvaluator.Factory( + source(), + ipEvaluatorSupplier, + prefixLengthV4EvaluatorSupplier, + prefixLengthV6EvaluatorSupplier, + context -> new BytesRef(new byte[16]) + ); + } + + @Evaluator(warnExceptions = IllegalArgumentException.class) + static BytesRef process( + BytesRef ip, + int prefixLengthV4, + int prefixLengthV6, + @Fixed(includeInToString = false, build = true) BytesRef scratch + ) { + if (prefixLengthV4 < 0 || prefixLengthV4 > 32) { + throw new IllegalArgumentException("Prefix length v4 must be in range [0, 32], found " + prefixLengthV4); + } + if (prefixLengthV6 < 0 || prefixLengthV6 > 128) { + throw new IllegalArgumentException("Prefix length v6 must be in range [0, 128], found " + prefixLengthV6); + } + + boolean isIpv4 = Arrays.compareUnsigned(ip.bytes, 0, IPV4_PREFIX.length, IPV4_PREFIX, 0, IPV4_PREFIX.length) == 0; + + if (isIpv4) { + makePrefix(ip, scratch, 12 + prefixLengthV4 / 8, prefixLengthV4 % 8); + } else { + makePrefix(ip, scratch, prefixLengthV6 / 8, prefixLengthV6 % 8); + } + + return scratch; + } + + private static void makePrefix(BytesRef ip, BytesRef scratch, int fullBytes, int remainingBits) { + // Copy the first full bytes + System.arraycopy(ip.bytes, ip.offset, scratch.bytes, 0, fullBytes); + + // Copy the last byte ignoring the trailing bits + if (remainingBits > 0) { + byte lastByteMask = (byte) (0xFF << (8 - remainingBits)); + scratch.bytes[fullBytes] = (byte) (ip.bytes[fullBytes] & lastByteMask); + } + + // Copy the last empty bytes + if (fullBytes < 16) { + Arrays.fill(scratch.bytes, fullBytes + 1, 16, (byte) 0); + } + } + + @Override + public DataType dataType() { + return DataTypes.IP; + } + + @Override + protected TypeResolution resolveType() { + if (childrenResolved() == false) { + return new TypeResolution("Unresolved children"); + } + + return isIPAndExact(ipField, sourceText(), FIRST).and( + isType(prefixLengthV4Field, dt -> dt == INTEGER, sourceText(), SECOND, "integer") + ).and(isType(prefixLengthV6Field, dt -> dt == INTEGER, sourceText(), THIRD, "integer")); + } + + @Override + public Expression replaceChildren(List newChildren) { + return new IpPrefix(source(), newChildren.get(0), newChildren.get(1), newChildren.get(2)); + } + + @Override + protected NodeInfo info() { + return NodeInfo.create(this, IpPrefix::new, ipField, prefixLengthV4Field, prefixLengthV6Field); + } +} diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/io/stream/PlanNamedTypes.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/io/stream/PlanNamedTypes.java index e846360258eb..67a6a5d4bb98 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/io/stream/PlanNamedTypes.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/io/stream/PlanNamedTypes.java @@ -98,6 +98,7 @@ import org.elasticsearch.xpack.esql.expression.function.scalar.date.DateParse; import org.elasticsearch.xpack.esql.expression.function.scalar.date.DateTrunc; import org.elasticsearch.xpack.esql.expression.function.scalar.date.Now; import org.elasticsearch.xpack.esql.expression.function.scalar.ip.CIDRMatch; +import org.elasticsearch.xpack.esql.expression.function.scalar.ip.IpPrefix; import org.elasticsearch.xpack.esql.expression.function.scalar.math.Abs; import org.elasticsearch.xpack.esql.expression.function.scalar.math.Acos; import org.elasticsearch.xpack.esql.expression.function.scalar.math.Asin; @@ -394,6 +395,7 @@ public final class PlanNamedTypes { of(ScalarFunction.class, DateTrunc.class, PlanNamedTypes::writeDateTrunc, PlanNamedTypes::readDateTrunc), of(ScalarFunction.class, E.class, PlanNamedTypes::writeNoArgScalar, PlanNamedTypes::readNoArgScalar), of(ScalarFunction.class, Greatest.class, PlanNamedTypes::writeVararg, PlanNamedTypes::readVarag), + of(ScalarFunction.class, IpPrefix.class, IpPrefix::writeTo, IpPrefix::readFrom), of(ScalarFunction.class, Least.class, PlanNamedTypes::writeVararg, PlanNamedTypes::readVarag), of(ScalarFunction.class, Log.class, PlanNamedTypes::writeLog, PlanNamedTypes::readLog), of(ScalarFunction.class, Now.class, PlanNamedTypes::writeNow, PlanNamedTypes::readNow), diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/ip/IpPrefixTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/ip/IpPrefixTests.java new file mode 100644 index 000000000000..e46eaea849bb --- /dev/null +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/ip/IpPrefixTests.java @@ -0,0 +1,116 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.esql.expression.function.scalar.ip; + +import com.carrotsearch.randomizedtesting.annotations.Name; +import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; + +import org.elasticsearch.common.network.InetAddresses; +import org.elasticsearch.common.network.NetworkAddress; +import org.elasticsearch.test.ESTestCase; +import org.elasticsearch.xpack.esql.core.expression.Expression; +import org.elasticsearch.xpack.esql.core.tree.Source; +import org.elasticsearch.xpack.esql.core.type.DataTypes; +import org.elasticsearch.xpack.esql.expression.function.AbstractFunctionTestCase; +import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier; +import org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter; + +import java.util.List; +import java.util.function.Supplier; + +import static org.hamcrest.Matchers.equalTo; + +public class IpPrefixTests extends AbstractFunctionTestCase { + public IpPrefixTests(@Name("TestCase") Supplier testCaseSupplier) { + this.testCase = testCaseSupplier.get(); + } + + @ParametersFactory + public static Iterable parameters() { + var suppliers = List.of( + // V4 + new TestCaseSupplier( + List.of(DataTypes.IP, DataTypes.INTEGER, DataTypes.INTEGER), + () -> new TestCaseSupplier.TestCase( + List.of( + new TestCaseSupplier.TypedData(EsqlDataTypeConverter.stringToIP("1.2.3.4"), DataTypes.IP, "ip"), + new TestCaseSupplier.TypedData(24, DataTypes.INTEGER, "prefixLengthV4"), + new TestCaseSupplier.TypedData(ESTestCase.randomIntBetween(0, 128), DataTypes.INTEGER, "prefixLengthV6") + ), + "IpPrefixEvaluator[ip=Attribute[channel=0], prefixLengthV4=Attribute[channel=1], prefixLengthV6=Attribute[channel=2]]", + DataTypes.IP, + equalTo(EsqlDataTypeConverter.stringToIP("1.2.3.0")) + ) + ), + new TestCaseSupplier(List.of(DataTypes.IP, DataTypes.INTEGER, DataTypes.INTEGER), () -> { + var randomIp = randomIp(true); + var randomPrefix = randomIntBetween(0, 32); + var cidrString = InetAddresses.toCidrString(randomIp, randomPrefix); + + var ipParameter = EsqlDataTypeConverter.stringToIP(NetworkAddress.format(randomIp)); + var expectedPrefix = EsqlDataTypeConverter.stringToIP( + NetworkAddress.format(InetAddresses.parseIpRangeFromCidr(cidrString).lowerBound()) + ); + + return new TestCaseSupplier.TestCase( + List.of( + new TestCaseSupplier.TypedData(ipParameter, DataTypes.IP, "ip"), + new TestCaseSupplier.TypedData(randomPrefix, DataTypes.INTEGER, "prefixLengthV4"), + new TestCaseSupplier.TypedData(ESTestCase.randomIntBetween(0, 128), DataTypes.INTEGER, "prefixLengthV6") + ), + "IpPrefixEvaluator[ip=Attribute[channel=0], prefixLengthV4=Attribute[channel=1], prefixLengthV6=Attribute[channel=2]]", + DataTypes.IP, + equalTo(expectedPrefix) + ); + }), + + // V6 + new TestCaseSupplier( + List.of(DataTypes.IP, DataTypes.INTEGER, DataTypes.INTEGER), + () -> new TestCaseSupplier.TestCase( + List.of( + new TestCaseSupplier.TypedData(EsqlDataTypeConverter.stringToIP("::ff"), DataTypes.IP, "ip"), + new TestCaseSupplier.TypedData(ESTestCase.randomIntBetween(0, 32), DataTypes.INTEGER, "prefixLengthV4"), + new TestCaseSupplier.TypedData(127, DataTypes.INTEGER, "prefixLengthV6") + ), + "IpPrefixEvaluator[ip=Attribute[channel=0], prefixLengthV4=Attribute[channel=1], prefixLengthV6=Attribute[channel=2]]", + DataTypes.IP, + equalTo(EsqlDataTypeConverter.stringToIP("::fe")) + ) + ), + new TestCaseSupplier(List.of(DataTypes.IP, DataTypes.INTEGER, DataTypes.INTEGER), () -> { + var randomIp = randomIp(false); + var randomPrefix = randomIntBetween(0, 128); + var cidrString = InetAddresses.toCidrString(randomIp, randomPrefix); + + var ipParameter = EsqlDataTypeConverter.stringToIP(NetworkAddress.format(randomIp)); + var expectedPrefix = EsqlDataTypeConverter.stringToIP( + NetworkAddress.format(InetAddresses.parseIpRangeFromCidr(cidrString).lowerBound()) + ); + + return new TestCaseSupplier.TestCase( + List.of( + new TestCaseSupplier.TypedData(ipParameter, DataTypes.IP, "ip"), + new TestCaseSupplier.TypedData(ESTestCase.randomIntBetween(0, 32), DataTypes.INTEGER, "prefixLengthV4"), + new TestCaseSupplier.TypedData(randomPrefix, DataTypes.INTEGER, "prefixLengthV6") + ), + "IpPrefixEvaluator[ip=Attribute[channel=0], prefixLengthV4=Attribute[channel=1], prefixLengthV6=Attribute[channel=2]]", + DataTypes.IP, + equalTo(expectedPrefix) + ); + }) + ); + + return parameterSuppliersFromTypedData(errorsForCasesWithoutExamples(anyNullIsNull(true, suppliers))); + } + + @Override + protected Expression build(Source source, List args) { + return new IpPrefix(source, args.get(0), args.get(1), args.size() == 3 ? args.get(2) : null); + } +}