Add Term/TermInSet query support for DATE_TIME field (#797)

* Add Term/TermInSet query support for DATE_TIME field * Unify doc value query building
Yelp · Jan 6, 2025 · 6caea92 · 6caea92
1 parent 3472026
commit 6caea92
Show file tree

Hide file tree

Showing 5 changed files with 410 additions and 17 deletions.
diff --git a/src/main/java/com/yelp/nrtsearch/server/field/DateTimeFieldDef.java b/src/main/java/com/yelp/nrtsearch/server/field/DateTimeFieldDef.java
@@ -20,6 +20,7 @@
 import com.yelp.nrtsearch.server.doc.LoadedDocValues;
 import com.yelp.nrtsearch.server.field.properties.RangeQueryable;
 import com.yelp.nrtsearch.server.field.properties.Sortable;
+import com.yelp.nrtsearch.server.field.properties.TermQueryable;
 import com.yelp.nrtsearch.server.grpc.*;
 import com.yelp.nrtsearch.server.grpc.Field;
 import java.io.IOException;
@@ -32,6 +33,7 @@
 import java.time.format.DateTimeParseException;
 import java.time.format.ResolverStyle;
 import java.time.temporal.ChronoField;
+import java.util.ArrayList;
 import java.util.List;
 import org.apache.lucene.document.*;
 import org.apache.lucene.facet.FacetField;
@@ -46,7 +48,7 @@
 
 /** Field class for 'DATE_TIME' field type. */
 public class DateTimeFieldDef extends IndexableFieldDef<Instant>
-    implements Sortable, RangeQueryable {
+    implements Sortable, RangeQueryable, TermQueryable {
   private static final String EPOCH_MILLIS = "epoch_millis";
   private static final String STRICT_DATE_OPTIONAL_TIME = "strict_date_optional_time";
 
@@ -140,6 +142,67 @@ public Query getRangeQuery(RangeQuery rangeQuery) {
     return new IndexOrDocValuesQuery(pointQuery, dvQuery);
   }
 
+  @Override
+  public void checkTermQueriesSupported() {
+    if (!isSearchable() && !hasDocValues()) {
+      throw new IllegalStateException(
+          "Field \""
+              + getName()
+              + "\" is not searchable or does not have doc values, which is required for TermQuery / TermInSetQuery");
+    }
+  }
+
+  @Override
+  public Query getTermQueryFromLongValue(long longValue) {
+    Query pointQuery = null;
+    Query dvQuery = null;
+    if (isSearchable()) {
+      pointQuery = LongPoint.newExactQuery(getName(), longValue);
+      if (!hasDocValues()) {
+        return pointQuery;
+      }
+    }
+    if (hasDocValues()) {
+      dvQuery = SortedNumericDocValuesField.newSlowExactQuery(getName(), longValue);
+      if (!isSearchable()) {
+        return dvQuery;
+      }
+    }
+    return new IndexOrDocValuesQuery(pointQuery, dvQuery);
+  }
+
+  @Override
+  public Query getTermInSetQueryFromLongValues(List<Long> longValues) {
+    Query pointQuery = null;
+    Query dvQuery = null;
+    if (isSearchable()) {
+      pointQuery = LongPoint.newSetQuery(getName(), longValues);
+      if (!hasDocValues()) {
+        return pointQuery;
+      }
+    }
+    if (hasDocValues()) {
+      long[] longValuesArray = longValues.stream().mapToLong(l -> l).toArray();
+      dvQuery = SortedNumericDocValuesField.newSlowSetQuery(getName(), longValuesArray);
+      if (!isSearchable()) {
+        return dvQuery;
+      }
+    }
+    return new IndexOrDocValuesQuery(pointQuery, dvQuery);
+  }
+
+  @Override
+  public Query getTermQueryFromTextValue(String textValue) {
+    return getTermQueryFromLongValue(getTimeToIndex(textValue));
+  }
+
+  @Override
+  public Query getTermInSetQueryFromTextValues(List<String> textValues) {
+    List<Long> longTerms = new ArrayList<>(textValues.size());
+    textValues.forEach((s) -> longTerms.add(getTimeToIndex(s)));
+    return getTermInSetQueryFromLongValues(longTerms);
+  }
+
   private long convertDateStringToMillis(String dateString) {
     if (dateTimeFormat.equals(EPOCH_MILLIS)) {
       return Long.parseLong(dateString);

diff --git a/src/main/java/com/yelp/nrtsearch/server/field/properties/TermQueryable.java b/src/main/java/com/yelp/nrtsearch/server/field/properties/TermQueryable.java
@@ -16,6 +16,7 @@
 package com.yelp.nrtsearch.server.field.properties;
 
 import com.yelp.nrtsearch.server.field.FieldDef;
+import com.yelp.nrtsearch.server.field.IndexableFieldDef;
 import com.yelp.nrtsearch.server.grpc.TermInSetQuery;
 import com.yelp.nrtsearch.server.grpc.TermQuery;
 import java.util.List;
@@ -260,4 +261,21 @@ default Query getTermInSetQueryFromLongValues(List<Long> longValues) {
   default Query getTermInSetQueryFromTextValues(List<String> textValues) {
     return null;
   }
+
+  /**
+   * Verify that this field supports term/term in set queries.
+   *
+   * @throws IllegalStateException if term queries are not supported
+   */
+  default void checkTermQueriesSupported() {
+    if (!(this instanceof IndexableFieldDef<?> indexableFieldDef)) {
+      throw new IllegalStateException("Instance is not an IndexableFieldDef");
+    }
+    if (!indexableFieldDef.isSearchable()) {
+      throw new IllegalStateException(
+          "Field "
+              + indexableFieldDef.getName()
+              + " is not searchable, which is required for TermQuery / TermInSetQuery");
+    }
+  }
 }
diff --git a/src/main/java/com/yelp/nrtsearch/server/query/QueryNodeMapper.java b/src/main/java/com/yelp/nrtsearch/server/query/QueryNodeMapper.java
@@ -294,34 +294,24 @@ private Query getTermQuery(com.yelp.nrtsearch.server.grpc.TermQuery termQuery, I
     String fieldName = termQuery.getField();
     FieldDef fieldDef = state.getFieldOrThrow(fieldName);
 
-    if (fieldDef instanceof TermQueryable) {
-      validateTermQueryIsSearchable(fieldDef);
-      return ((TermQueryable) fieldDef).getTermQuery(termQuery);
+    if (fieldDef instanceof TermQueryable termQueryable) {
+      termQueryable.checkTermQueriesSupported();
+      return termQueryable.getTermQuery(termQuery);
     }
 
     String message =
         "Unable to create TermQuery: %s, field type: %s is not supported for TermQuery";
     throw new IllegalArgumentException(String.format(message, termQuery, fieldDef.getType()));
   }
 
-  private void validateTermQueryIsSearchable(FieldDef fieldDef) {
-    if (fieldDef instanceof IndexableFieldDef
-        && !((IndexableFieldDef<?>) fieldDef).isSearchable()) {
-      throw new IllegalStateException(
-          "Field "
-              + fieldDef.getName()
-              + " is not searchable, which is required for TermQuery / TermInSetQuery");
-    }
-  }
-
   private Query getTermInSetQuery(
       com.yelp.nrtsearch.server.grpc.TermInSetQuery termInSetQuery, IndexState state) {
     String fieldName = termInSetQuery.getField();
     FieldDef fieldDef = state.getFieldOrThrow(fieldName);
 
-    if (fieldDef instanceof TermQueryable) {
-      validateTermQueryIsSearchable(fieldDef);
-      return ((TermQueryable) fieldDef).getTermInSetQuery(termInSetQuery);
+    if (fieldDef instanceof TermQueryable termQueryable) {
+      termQueryable.checkTermQueriesSupported();
+      return termQueryable.getTermInSetQuery(termInSetQuery);
     }
 
     String message =