Quick hack via MultiRangeQuery from Lucene sandbox

Signed-off-by: mikhail-khludnev <mikhail_khludnev@rntgroup.com>
opensearch-project · Oct 19, 2024 · c9e2bd1 · c9e2bd1
1 parent 853a344
commit c9e2bd1
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 4 deletions.
diff --git a/server/src/main/java/org/opensearch/index/mapper/IpFieldMapper.java b/server/src/main/java/org/opensearch/index/mapper/IpFieldMapper.java
@@ -36,6 +36,7 @@
 import org.apache.lucene.document.SortedSetDocValuesField;
 import org.apache.lucene.document.StoredField;
 import org.apache.lucene.index.SortedSetDocValues;
+import org.apache.lucene.sandbox.search.MultiRangeQuery;
 import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.BooleanQuery;
 import org.apache.lucene.search.ConstantScoreQuery;
@@ -273,6 +274,8 @@ public Query termsQuery(List<?> values, QueryShardContext context) {
             failIfNotIndexedAndNoDocValues();
             List<InetAddress> concreteIPs = new ArrayList<>();
             List<Query> ranges = new ArrayList<>();
+            IpMultiRangeQueryBuilder multiRange = new IpMultiRangeQueryBuilder(name());
+            boolean multiRangeIsEmpty = true;
             for (final Object value : values) {
                 if (value instanceof InetAddress) {
                     concreteIPs.add((InetAddress) value);
@@ -281,14 +284,22 @@ public Query termsQuery(List<?> values, QueryShardContext context) {
                     if (strVal.contains("/")) {
                         // the `terms` query contains some prefix queries, so we cannot create a set query
                         // and need to fall back to a disjunction of `term` queries
-                        Query query = termQuery(strVal, context);
+                        // Query query = termQuery(strVal, context);
+                        final Tuple<InetAddress, Integer> cidr = InetAddresses.parseCidr(strVal);
+                        PointRangeQuery query = (PointRangeQuery) InetAddressPoint.newPrefixQuery(name(), cidr.v1(), cidr.v2());
+
                         // would be great to have union on ranges over bare points
-                        ranges.add(query);
+                        // ranges.add(query);
+                        multiRange.add(query.getLowerPoint(), query.getUpperPoint());
+                        multiRangeIsEmpty = false;
                     } else {
                         concreteIPs.add(InetAddresses.forString(strVal));
                     }
                 }
             }
+            if (!multiRangeIsEmpty) {
+                ranges.add(multiRange.build());
+            }
             if (!concreteIPs.isEmpty()) {
                 Supplier<Query> pointsQuery;
                 pointsQuery = () -> concreteIPs.size() == 1
@@ -470,6 +481,27 @@ public DocValueFormat docValueFormat(@Nullable String format, ZoneId timeZone) {
         }
     }
 
+    public static class IpMultiRangeQueryBuilder extends MultiRangeQuery.Builder {
+        public IpMultiRangeQueryBuilder(String field) {
+            super(field, InetAddressPoint.BYTES, 1);
+        }
+
+        public IpMultiRangeQueryBuilder add(InetAddress lower, InetAddress upper) {
+            add(new MultiRangeQuery.RangeClause(InetAddressPoint.encode(lower), InetAddressPoint.encode(upper)));
+            return this;
+        }
+
+        @Override
+        public MultiRangeQuery build() {
+            return new MultiRangeQuery(field, numDims, bytesPerDim, clauses) {
+                @Override
+                protected String toString(int dimension, byte[] value) {
+                    return InetAddressPoint.decode(value).getHostAddress();
+                }
+            };
+        }
+    }
+
     private final boolean indexed;
     private final boolean hasDocValues;
     private final boolean stored;

diff --git a/server/src/test/java/org/opensearch/search/SearchIpFieldTermsTests.java b/server/src/test/java/org/opensearch/search/SearchIpFieldTermsTests.java
@@ -42,7 +42,7 @@ public void testMassive() throws Exception {
         int cidrs = 0;
         int ips = 0;
         List<String> toQuery = new ArrayList<>();
-        for (int i = 0; ips <= 1024 && i < 1000000; i++) {
+        for (int i = 0; ips <= 10240 && cidrs <= 1024 && i < 1000000; i++) {
             final String ip;
             final int prefix;
             if (IPv4_ONLY) {
@@ -56,7 +56,7 @@ public void testMassive() throws Exception {
             bulkRequestBuilder.add(client().prepareIndex(defaultIndexName).setSource(Map.of("addr", ip)));
 
             final String termToQuery;
-            if (cidrs < 1024 - 1 && random().nextBoolean()) {
+            if (random().nextBoolean()) {
                 termToQuery = ip + "/" + prefix;
                 cidrs++;
             } else {