Add range query support for _ID and ATOM fields (#804)

Yelp · Jan 7, 2025 · bb431d0 · bb431d0
1 parent 4b4dbd8
commit bb431d0
Show file tree

Hide file tree

Showing 5 changed files with 446 additions and 49 deletions.
diff --git a/src/main/java/com/yelp/nrtsearch/server/field/AtomFieldDef.java b/src/main/java/com/yelp/nrtsearch/server/field/AtomFieldDef.java
@@ -17,19 +17,25 @@
 
 import static com.yelp.nrtsearch.server.analysis.AnalyzerCreator.hasAnalyzer;
 
+import com.yelp.nrtsearch.server.field.properties.RangeQueryable;
 import com.yelp.nrtsearch.server.field.properties.Sortable;
 import com.yelp.nrtsearch.server.grpc.Field;
+import com.yelp.nrtsearch.server.grpc.RangeQuery;
 import com.yelp.nrtsearch.server.grpc.SortType;
 import org.apache.lucene.analysis.Analyzer;
 import org.apache.lucene.analysis.core.KeywordAnalyzer;
 import org.apache.lucene.document.FieldType;
+import org.apache.lucene.document.SortedSetDocValuesField;
+import org.apache.lucene.index.DocValuesType;
 import org.apache.lucene.index.IndexOptions;
+import org.apache.lucene.search.Query;
 import org.apache.lucene.search.SortField;
 import org.apache.lucene.search.SortedSetSortField;
+import org.apache.lucene.search.TermRangeQuery;
 import org.apache.lucene.util.BytesRef;
 
 /** Field class for 'ATOM' field type. Uses {@link KeywordAnalyzer} for text analysis. */
-public class AtomFieldDef extends TextBaseFieldDef implements Sortable {
+public class AtomFieldDef extends TextBaseFieldDef implements Sortable, RangeQueryable {
   private static final Analyzer keywordAnalyzer = new KeywordAnalyzer();
 
   public AtomFieldDef(
@@ -98,4 +104,32 @@ public SortField getSortField(SortType type) {
     }
     return sortField;
   }
+
+  @Override
+  public Query getRangeQuery(RangeQuery rangeQuery) {
+    verifySearchableOrDocValues("Range query");
+    BytesRef lowerTerm =
+        rangeQuery.getLower().isEmpty() ? null : new BytesRef(rangeQuery.getLower());
+    BytesRef upperTerm =
+        rangeQuery.getUpper().isEmpty() ? null : new BytesRef(rangeQuery.getUpper());
+    if (isSearchable()) {
+      return new TermRangeQuery(
+          getName(),
+          lowerTerm,
+          upperTerm,
+          !rangeQuery.getLowerExclusive(),
+          !rangeQuery.getUpperExclusive());
+    } else if (hasDocValues()
+        && (docValuesType == DocValuesType.SORTED || docValuesType == DocValuesType.SORTED_SET)) {
+      return SortedSetDocValuesField.newSlowRangeQuery(
+          getName(),
+          lowerTerm,
+          upperTerm,
+          !rangeQuery.getLowerExclusive(),
+          !rangeQuery.getUpperExclusive());
+    } else {
+      throw new IllegalStateException(
+          "Only SORTED or SORTED_SET doc values are supported for range queries: " + getName());
+    }
+  }
 }
diff --git a/src/main/java/com/yelp/nrtsearch/server/field/IdFieldDef.java b/src/main/java/com/yelp/nrtsearch/server/field/IdFieldDef.java
@@ -16,8 +16,10 @@
 package com.yelp.nrtsearch.server.field;
 
 import com.yelp.nrtsearch.server.doc.LoadedDocValues;
+import com.yelp.nrtsearch.server.field.properties.RangeQueryable;
 import com.yelp.nrtsearch.server.field.properties.TermQueryable;
 import com.yelp.nrtsearch.server.grpc.Field;
+import com.yelp.nrtsearch.server.grpc.RangeQuery;
 import com.yelp.nrtsearch.server.grpc.SearchResponse;
 import java.io.IOException;
 import java.util.List;
@@ -33,10 +35,11 @@
 import org.apache.lucene.index.SortedDocValues;
 import org.apache.lucene.index.Term;
 import org.apache.lucene.search.Query;
+import org.apache.lucene.search.TermRangeQuery;
 import org.apache.lucene.util.BytesRef;
 
 /** Field class for defining '_ID' fields which are used to update documents */
-public class IdFieldDef extends IndexableFieldDef<String> implements TermQueryable {
+public class IdFieldDef extends IndexableFieldDef<String> implements TermQueryable, RangeQueryable {
 
   protected IdFieldDef(
       String name, Field requestField, FieldDefCreator.FieldDefCreatorContext context) {
@@ -169,4 +172,19 @@ public Query getTermInSetQueryFromTextValues(List<String> textValues) {
     List<BytesRef> textTerms = textValues.stream().map(BytesRef::new).collect(Collectors.toList());
     return new org.apache.lucene.search.TermInSetQuery(getName(), textTerms);
   }
+
+  @Override
+  public Query getRangeQuery(RangeQuery rangeQuery) {
+    // _ID fields are always searchable
+    BytesRef lowerTerm =
+        rangeQuery.getLower().isEmpty() ? null : new BytesRef(rangeQuery.getLower());
+    BytesRef upperTerm =
+        rangeQuery.getUpper().isEmpty() ? null : new BytesRef(rangeQuery.getUpper());
+    return new TermRangeQuery(
+        getName(),
+        lowerTerm,
+        upperTerm,
+        !rangeQuery.getLowerExclusive(),
+        !rangeQuery.getUpperExclusive());
+  }
 }