support rangeQuery and regexpQuery in constant_keyword field type

Signed-off-by: kkewwei <[email protected]>
opensearch-project · Jul 11, 2024 · 3cd9c3e · 3cd9c3e
1 parent dfb8449
commit 3cd9c3e
Show file tree

Hide file tree

Showing 3 changed files with 121 additions and 1 deletion.
diff --git a/server/src/main/java/org/opensearch/index/mapper/ConstantFieldType.java b/server/src/main/java/org/opensearch/index/mapper/ConstantFieldType.java
@@ -76,7 +76,7 @@ public final boolean isAggregatable() {
      */
     protected abstract boolean matches(String pattern, boolean caseInsensitive, QueryShardContext context);
 
-    private static String valueToString(Object value) {
+    protected static String valueToString(Object value) {
         return value instanceof BytesRef ? ((BytesRef) value).utf8ToString() : value.toString();
     }
 

diff --git a/server/src/main/java/org/opensearch/index/mapper/ConstantKeywordFieldMapper.java b/server/src/main/java/org/opensearch/index/mapper/ConstantKeywordFieldMapper.java
@@ -9,17 +9,29 @@
 package org.opensearch.index.mapper;
 
 import org.apache.lucene.search.MatchAllDocsQuery;
+import org.apache.lucene.search.MatchNoDocsQuery;
+import org.apache.lucene.search.MultiTermQuery;
 import org.apache.lucene.search.Query;
+import org.apache.lucene.search.RegexpQuery;
+import org.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.automaton.Automaton;
+import org.apache.lucene.util.automaton.ByteRunAutomaton;
+import org.apache.lucene.util.automaton.RegExp;
 import org.opensearch.OpenSearchParseException;
+import org.opensearch.common.Nullable;
 import org.opensearch.common.annotation.PublicApi;
+import org.opensearch.common.geo.ShapeRelation;
+import org.opensearch.common.lucene.BytesRefs;
 import org.opensearch.common.regex.Regex;
+import org.opensearch.common.time.DateMathParser;
 import org.opensearch.index.fielddata.IndexFieldData;
 import org.opensearch.index.fielddata.plain.ConstantIndexFieldData;
 import org.opensearch.index.query.QueryShardContext;
 import org.opensearch.search.aggregations.support.CoreValuesSourceType;
 import org.opensearch.search.lookup.SearchLookup;
 
 import java.io.IOException;
+import java.time.ZoneId;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
@@ -122,6 +134,60 @@ public Query existsQuery(QueryShardContext context) {
             return new MatchAllDocsQuery();
         }
 
+        @Override
+        public Query rangeQuery(
+            Object lowerTerm,
+            Object upperTerm,
+            boolean includeLower,
+            boolean includeUpper,
+            ShapeRelation relation,
+            ZoneId timeZone,
+            DateMathParser parser,
+            QueryShardContext context
+        ) {
+            if (lowerTerm != null) {
+                lowerTerm = valueToString(lowerTerm);
+            }
+            if (upperTerm != null) {
+                upperTerm = valueToString(upperTerm);
+            }
+
+            if (lowerTerm != null && upperTerm != null && ((String) lowerTerm).compareTo((String) upperTerm) > 0) {
+                return new MatchNoDocsQuery();
+            }
+
+            if (lowerTerm != null && ((String) lowerTerm).compareTo(value) > (includeLower ? 0 : -1)) {
+                return new MatchNoDocsQuery();
+            }
+
+            if (upperTerm != null && ((String) upperTerm).compareTo(value) < (includeUpper ? 0 : 1)) {
+                return new MatchNoDocsQuery();
+            }
+            return new MatchAllDocsQuery();
+        }
+
+        @Override
+        public Query regexpQuery(
+            String value,
+            int syntaxFlags,
+            int matchFlags,
+            int maxDeterminizedStates,
+            @Nullable MultiTermQuery.RewriteMethod method,
+            QueryShardContext context
+        ) {
+            Automaton automaton = new RegExp(value, syntaxFlags, matchFlags).toAutomaton(
+                RegexpQuery.DEFAULT_PROVIDER,
+                maxDeterminizedStates
+            );
+            ByteRunAutomaton byteRunAutomaton = new ByteRunAutomaton(automaton);
+            BytesRef valueBytes = BytesRefs.toBytesRef(this.value);
+            if (byteRunAutomaton.run(valueBytes.bytes, valueBytes.offset, valueBytes.length)) {
+                return new MatchAllDocsQuery();
+            } else {
+                return new MatchNoDocsQuery();
+            }
+        }
+
         @Override
         public IndexFieldData.Builder fielddataBuilder(String fullyQualifiedIndexName, Supplier<SearchLookup> searchLookup) {
             return new ConstantIndexFieldData.Builder(fullyQualifiedIndexName, name(), CoreValuesSourceType.BYTES);

diff --git a/server/src/test/java/org/opensearch/index/mapper/ConstantKeywordFieldTypeTests.java b/server/src/test/java/org/opensearch/index/mapper/ConstantKeywordFieldTypeTests.java
@@ -10,6 +10,8 @@
 
 import org.apache.lucene.search.MatchAllDocsQuery;
 import org.apache.lucene.search.MatchNoDocsQuery;
+import org.apache.lucene.search.MultiTermQuery;
+import org.apache.lucene.search.Query;
 import org.opensearch.Version;
 import org.opensearch.cluster.metadata.IndexMetadata;
 import org.opensearch.common.regex.Regex;
@@ -61,6 +63,58 @@ public void testExistsQuery() {
         assertEquals(new MatchAllDocsQuery(), ft.existsQuery(createContext()));
     }
 
+    public void testRangeQuery() {
+        Query actual = ft.rangeQuery("default", null, true, false, null, null, null, MOCK_QSC);
+        assertEquals(new MatchAllDocsQuery(), actual);
+
+        actual = ft.rangeQuery("default", null, false, false, null, null, null, MOCK_QSC);
+        assertEquals(new MatchNoDocsQuery(), actual);
+
+        actual = ft.rangeQuery(null, "default", true, true, null, null, null, MOCK_QSC);
+        assertEquals(new MatchAllDocsQuery(), actual);
+
+        actual = ft.rangeQuery(null, "default", false, false, null, null, null, MOCK_QSC);
+        assertEquals(new MatchNoDocsQuery(), actual);
+
+        actual = ft.rangeQuery("default", "default", false, true, null, null, null, MOCK_QSC);
+        assertEquals(new MatchNoDocsQuery(), actual);
+
+        actual = ft.rangeQuery("default", "default", true, false, null, null, null, MOCK_QSC);
+        assertEquals(new MatchNoDocsQuery(), actual);
+
+        actual = ft.rangeQuery(null, null, false, false, null, null, null, MOCK_QSC);
+        assertEquals(new MatchAllDocsQuery(), actual);
+
+        actual = ft.rangeQuery("default", "default", true, true, null, null, null, MOCK_QSC);
+        assertEquals(new MatchAllDocsQuery(), actual);
+
+        actual = ft.rangeQuery("defaul", "default1", true, true, null, null, null, MOCK_QSC);
+        assertEquals(new MatchAllDocsQuery(), actual);
+    }
+
+    public void testRegexpQuery() {
+        final ConstantKeywordFieldMapper.ConstantKeywordFieldType ft = new ConstantKeywordFieldMapper.ConstantKeywordFieldType(
+            "field",
+            "d3efault"
+        );
+        // test .*
+        Query query = ft.regexpQuery("d.*", 0, 0, 10, MultiTermQuery.CONSTANT_SCORE_BLENDED_REWRITE, MOCK_QSC);
+        assertEquals(new MatchAllDocsQuery(), query);
+        // test \d and ?
+        query = ft.regexpQuery("d\\defau[a-z]?t", 0, 0, 10, MultiTermQuery.CONSTANT_SCORE_BLENDED_REWRITE, MOCK_QSC);
+        assertEquals(new MatchAllDocsQuery(), query);
+
+        // test \d and ?
+        query = ft.regexpQuery("d\\defa[a-z]?t", 0, 0, 10, MultiTermQuery.CONSTANT_SCORE_BLENDED_REWRITE, MOCK_QSC);
+        assertEquals(new MatchNoDocsQuery(), query);
+        // \w{m,n}
+        query = ft.regexpQuery("d3efa[a-z]{3,3}", 0, 0, 10, MultiTermQuery.CONSTANT_SCORE_BLENDED_REWRITE, MOCK_QSC);
+        assertEquals(new MatchAllDocsQuery(), query);
+        // \w{m,n}
+        query = ft.regexpQuery("d3efa[a-z]{4,4}", 0, 0, 10, MultiTermQuery.CONSTANT_SCORE_BLENDED_REWRITE, MOCK_QSC);
+        assertEquals(new MatchNoDocsQuery(), query);
+    }
+
     private QueryShardContext createContext() {
         IndexMetadata indexMetadata = IndexMetadata.builder("index")
             .settings(Settings.builder().put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT))