add a new ID field

slackhq · Feb 6, 2024 · 627410e · 627410e
1 parent 452c461
commit 627410e
Show file tree

Hide file tree

Showing 5 changed files with 55 additions and 3 deletions.
diff --git a/kaldb/src/main/java/com/slack/kaldb/logstore/opensearch/OpenSearchAdapter.java b/kaldb/src/main/java/com/slack/kaldb/logstore/opensearch/OpenSearchAdapter.java
@@ -220,6 +220,8 @@ public void reloadSchema() {
           tryRegisterField(mapperService, entry.getValue().name, b -> b.field("type", "text"));
         } else if (entry.getValue().fieldType == FieldType.STRING) {
           tryRegisterField(mapperService, entry.getValue().name, b -> b.field("type", "keyword"));
+        } else if (entry.getValue().fieldType == FieldType.ID) {
+          tryRegisterField(mapperService, entry.getValue().name, b -> b.field("type", "keyword"));
         } else if (entry.getValue().fieldType == FieldType.INTEGER) {
           tryRegisterField(mapperService, entry.getValue().name, b -> b.field("type", "integer"));
         } else if (entry.getValue().fieldType == FieldType.LONG) {

diff --git a/kaldb/src/main/java/com/slack/kaldb/logstore/schema/SchemaAwareLogDocumentBuilderImpl.java b/kaldb/src/main/java/com/slack/kaldb/logstore/schema/SchemaAwareLogDocumentBuilderImpl.java
@@ -59,7 +59,6 @@ public static ImmutableMap<String, LuceneFieldDef> getDefaultLuceneFieldDefiniti
     }
 
     String[] fieldsAsString = {
-      LogMessage.SystemField.ID.fieldName,
       LogMessage.SystemField.INDEX.fieldName,
       LogMessage.ReservedField.TYPE.fieldName,
       LogMessage.ReservedField.HOSTNAME.fieldName,
@@ -78,6 +77,14 @@ public static ImmutableMap<String, LuceneFieldDef> getDefaultLuceneFieldDefiniti
           fieldName, new LuceneFieldDef(fieldName, FieldType.STRING.name, false, true, true));
     }
 
+    String[] fieldsAsIds = {
+      LogMessage.SystemField.ID.fieldName,
+    };
+    for (String fieldName : fieldsAsIds) {
+      fieldDefBuilder.put(
+          fieldName, new LuceneFieldDef(fieldName, FieldType.ID.name, false, true, true));
+    }
+
     String[] fieldsAsLong = {
       LogMessage.ReservedField.DURATION_MS.fieldName,
       LogMessage.SystemField.TIME_SINCE_EPOCH.fieldName,

diff --git a/kaldb/src/main/java/com/slack/kaldb/metadata/schema/FieldType.java b/kaldb/src/main/java/com/slack/kaldb/metadata/schema/FieldType.java
@@ -28,6 +28,7 @@
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.TermQuery;
 import org.apache.lucene.util.BytesRef;
+import org.opensearch.index.mapper.Uid;
 
 /** The FieldType enum describes the types of fields in a chunk. */
 public enum FieldType {
@@ -81,6 +82,32 @@ public Analyzer getAnalyzer(boolean quoted) {
       return KEYWORD_ANALYZER;
     }
   },
+  ID("id") {
+    @Override
+    public void addField(Document doc, String name, Object value, LuceneFieldDef fieldDef) {
+      BytesRef id = Uid.encodeId((String) value);
+      if (fieldDef.isIndexed) {
+        doc.add(new StringField(name, id, getStoreEnum(fieldDef.isStored)));
+      }
+      if (fieldDef.isStored) {
+        doc.add(new StoredField(name, (String) value));
+      }
+      if (fieldDef.storeDocValue) {
+        doc.add(new SortedDocValuesField(name, id));
+      }
+    }
+
+    @Override
+    public Query termQuery(String field, String queryText, Analyzer analyzer) {
+      final Term term = new Term(field, queryText);
+      return new TermQuery(term);
+    }
+
+    @Override
+    public Analyzer getAnalyzer(boolean quoted) {
+      return KEYWORD_ANALYZER;
+    }
+  },
   INTEGER("integer") {
     @Override
     public void addField(Document doc, String name, Object v, LuceneFieldDef fieldDef) {
@@ -373,7 +400,7 @@ private static Field.Store getStoreEnum(boolean isStored) {
   // Aliased Field Types are FieldTypes that can be considered as same type from a field conflict
   // detection perspective
   public static final List<Set<FieldType>> ALIASED_FIELD_TYPES =
-      ImmutableList.of(ImmutableSet.of(FieldType.STRING, FieldType.TEXT));
+      ImmutableList.of(ImmutableSet.of(FieldType.STRING, FieldType.TEXT, FieldType.ID));
 
   public static boolean areTypeAliasedFieldTypes(FieldType type1, FieldType type2) {
     for (Set<FieldType> s : ALIASED_FIELD_TYPES) {

diff --git a/kaldb/src/test/java/com/slack/kaldb/logstore/schema/DropPolicyTest.java b/kaldb/src/test/java/com/slack/kaldb/logstore/schema/DropPolicyTest.java
@@ -50,7 +50,7 @@ public void testBasicDocumentCreation() throws IOException {
     assertThat(MetricsUtil.getCount(CONVERT_FIELD_VALUE_COUNTER, meterRegistry)).isZero();
     assertThat(MetricsUtil.getCount(CONVERT_AND_DUPLICATE_FIELD_COUNTER, meterRegistry)).isZero();
     // Only string fields have doc values not text fields.
-    assertThat(docBuilder.getSchema().get("_id").fieldType.name).isEqualTo(FieldType.STRING.name);
+    assertThat(docBuilder.getSchema().get("_id").fieldType.name).isEqualTo(FieldType.ID.name);
     assertThat(
             testDocument.getFields().stream()
                 .filter(

diff --git a/kaldb/src/test/java/com/slack/kaldb/logstore/search/LogIndexSearcherImplTest.java b/kaldb/src/test/java/com/slack/kaldb/logstore/search/LogIndexSearcherImplTest.java
@@ -1725,4 +1725,20 @@ public void testConcurrentSearches() throws InterruptedException {
     assertThat(searchFailures.get()).isEqualTo(0);
     assertThat(successfulRuns.get()).isEqualTo(200);
   }
+
+  @Test
+  public void testSearchById() {
+    Instant time = Instant.ofEpochSecond(1593365471);
+    loadTestData(time);
+    SearchResult<LogMessage> index =
+        strictLogStore.logSearcher.search(
+            TEST_DATASET_NAME,
+            "_id:1",
+            time.toEpochMilli(),
+            time.plusSeconds(2).toEpochMilli(),
+            10,
+            new DateHistogramAggBuilder(
+                "1", LogMessage.SystemField.TIME_SINCE_EPOCH.fieldName, "1s"));
+    assertThat(index.hits.size()).isEqualTo(1);
+  }
 }