Complete keyword changes

Signed-off-by: Bharathwaj G <[email protected]>
opensearch-project · Oct 11, 2024 · 3571ec8 · 3571ec8
1 parent bc6fe9b
commit 3571ec8
Show file tree

Hide file tree

Showing 32 changed files with 749 additions and 168 deletions.
diff --git a/server/src/main/java/org/apache/lucene/index/DocValuesWriterWrapper.java b/server/src/main/java/org/apache/lucene/index/DocValuesWriterWrapper.java
@@ -0,0 +1,18 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.apache.lucene.index;
+
+import org.apache.lucene.search.DocIdSetIterator;
+
+/**
+ * Base wrapper class for DocValuesWriter.
+ */
+public abstract class DocValuesWriterWrapper<T extends DocIdSetIterator> {
+    public abstract T getDocValues();
+}
diff --git a/server/src/main/java/org/apache/lucene/index/SortedNumericDocValuesWriterWrapper.java b/server/src/main/java/org/apache/lucene/index/SortedNumericDocValuesWriterWrapper.java
@@ -18,7 +18,7 @@
  *
  * @opensearch.experimental
  */
-public class SortedNumericDocValuesWriterWrapper {
+public class SortedNumericDocValuesWriterWrapper extends DocValuesWriterWrapper<SortedNumericDocValues> {
 
     private final SortedNumericDocValuesWriter sortedNumericDocValuesWriter;
 
@@ -47,6 +47,7 @@ public void addValue(int docID, long value) {
      *
      * @return the {@link SortedNumericDocValues} instance
      */
+    @Override
     public SortedNumericDocValues getDocValues() {
         return sortedNumericDocValuesWriter.getDocValues();
     }

diff --git a/server/src/main/java/org/apache/lucene/index/SortedSetDocValuesWriterWrapper.java b/server/src/main/java/org/apache/lucene/index/SortedSetDocValuesWriterWrapper.java
@@ -0,0 +1,58 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.apache.lucene.index;
+
+import org.apache.lucene.util.ByteBlockPool;
+import org.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.Counter;
+
+/**
+ * A wrapper class for writing sorted set doc values.
+ * <p>
+ * This class provides a convenient way to add sorted set doc values to a field
+ * and retrieve the corresponding {@link SortedSetDocValues} instance.
+ *
+ * @opensearch.experimental
+ */
+public class SortedSetDocValuesWriterWrapper extends DocValuesWriterWrapper<SortedSetDocValues> {
+
+    private final SortedSetDocValuesWriter sortedSetDocValuesWriterWrapper;
+
+    /**
+     * Sole constructor. Constructs a new {@link SortedSetDocValuesWriterWrapper} instance.
+     *
+     * @param fieldInfo the field information for the field being written
+     * @param counter a counter for tracking memory usage
+     * @param byteBlockPool a byte block pool for allocating byte blocks
+     * @see SortedSetDocValuesWriter
+     */
+    public SortedSetDocValuesWriterWrapper(FieldInfo fieldInfo, Counter counter, ByteBlockPool byteBlockPool) {
+        sortedSetDocValuesWriterWrapper = new SortedSetDocValuesWriter(fieldInfo, counter, byteBlockPool);
+    }
+
+    /**
+     * Adds a bytes ref value to the sorted set doc values for the specified document.
+     *
+     * @param docID the document ID
+     * @param value the value to add
+     */
+    public void addValue(int docID, BytesRef value) {
+        sortedSetDocValuesWriterWrapper.addValue(docID, value);
+    }
+
+    /**
+     * Returns the {@link SortedSetDocValues} instance containing the sorted numeric doc values
+     *
+     * @return the {@link SortedSetDocValues} instance
+     */
+    @Override
+    public SortedSetDocValues getDocValues() {
+        return sortedSetDocValuesWriterWrapper.getDocValues();
+    }
+}
diff --git a/...ain/java/org/opensearch/index/codec/composite/composite99/Composite99DocValuesReader.java b/...ain/java/org/opensearch/index/codec/composite/composite99/Composite99DocValuesReader.java
@@ -15,6 +15,7 @@
 import org.apache.lucene.index.BinaryDocValues;
 import org.apache.lucene.index.CorruptIndexException;
 import org.apache.lucene.index.DocValues;
+import org.apache.lucene.index.DocValuesType;
 import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.FieldInfos;
 import org.apache.lucene.index.IndexFileNames;
@@ -40,6 +41,7 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
@@ -111,7 +113,7 @@ public Composite99DocValuesReader(DocValuesProducer producer, SegmentReadState r
                     readState.segmentInfo.getId(),
                     readState.segmentSuffix
                 );
-
+                Map<String, DocValuesType> dimensionFieldTypeMap = new HashMap<>();
                 while (true) {
 
                     // validate magic marker
@@ -155,13 +157,15 @@ public Composite99DocValuesReader(DocValuesProducer producer, SegmentReadState r
                             compositeIndexInputMap.put(compositeFieldName, starTreeIndexInput);
                             compositeIndexMetadataMap.put(compositeFieldName, starTreeMetadata);
 
-                            List<String> dimensionFields = starTreeMetadata.getDimensionFields();
-
+                            Map<String, DocValuesType> dimensionFieldToDocValuesMap = starTreeMetadata.getDimensionFields();
                             // generating star tree unique fields (fully qualified name for dimension and metrics)
-                            for (String dimensions : dimensionFields) {
-                                fields.add(fullyQualifiedFieldNameForStarTreeDimensionsDocValues(compositeFieldName, dimensions));
+                            for (Map.Entry<String, DocValuesType> dimensionEntry : dimensionFieldToDocValuesMap.entrySet()) {
+                                String dimName = fullyQualifiedFieldNameForStarTreeDimensionsDocValues(
+                                    compositeFieldName,
+                                    dimensionEntry.getKey()
+                                );
+                                fields.add(dimName);
                             }
-
                             // adding metric fields
                             for (Metric metric : starTreeMetadata.getMetrics()) {
                                 for (MetricStat metricStat : metric.getBaseMetrics()) {
@@ -184,7 +188,7 @@ public Composite99DocValuesReader(DocValuesProducer producer, SegmentReadState r
 
                 // populates the dummy list of field infos to fetch doc id set iterators for respective fields.
                 // the dummy field info is used to fetch the doc id set iterators for respective fields based on field name
-                FieldInfos fieldInfos = new FieldInfos(getFieldInfoList(fields));
+                FieldInfos fieldInfos = new FieldInfos(getFieldInfoList(fields, dimensionFieldTypeMap));
                 this.readState = new SegmentReadState(readState.directory, readState.segmentInfo, fieldInfos, readState.context);
 
                 // initialize star-tree doc values producer
@@ -298,4 +302,8 @@ public static SortedNumericDocValues getSortedNumericDocValues(SortedNumericDocV
         return sortedNumeric == null ? DocValues.emptySortedNumeric() : sortedNumeric;
     }
 
+    public static SortedSetDocValues getSortedSetDocValues(SortedSetDocValues sortedSetDv) {
+        return sortedSetDv == null ? DocValues.emptySortedSet() : sortedSetDv;
+    }
+
 }
diff --git a/...ain/java/org/opensearch/index/codec/composite/composite99/Composite99DocValuesWriter.java b/...ain/java/org/opensearch/index/codec/composite/composite99/Composite99DocValuesWriter.java
@@ -73,6 +73,7 @@ public class Composite99DocValuesWriter extends DocValuesConsumer {
     private final AtomicInteger fieldNumberAcrossCompositeFields;
 
     private final Map<String, DocValuesProducer> fieldProducerMap = new HashMap<>();
+    private final Map<String, SortedSetDocValues> fieldDocIdSetIteratorMap = new HashMap<>();
 
     public Composite99DocValuesWriter(DocValuesConsumer delegate, SegmentWriteState segmentWriteState, MapperService mapperService)
         throws IOException {
@@ -191,6 +192,11 @@ public void addSortedSetField(FieldInfo field, DocValuesProducer valuesProducer)
         if (mergeState.get() == null && segmentHasCompositeFields) {
             createCompositeIndicesIfPossible(valuesProducer, field);
         }
+        if (mergeState.get() != null) {
+            if (compositeFieldSet.contains(field.name)) {
+                fieldDocIdSetIteratorMap.put(field.name, valuesProducer.getSortedSet(field));
+            }
+        }
     }
 
     @Override
@@ -339,7 +345,14 @@ private void mergeStarTreeFields(MergeState mergeState) throws IOException {
             }
         }
         try (StarTreesBuilder starTreesBuilder = new StarTreesBuilder(state, mapperService, fieldNumberAcrossCompositeFields)) {
-            starTreesBuilder.buildDuringMerge(metaOut, dataOut, starTreeSubsPerField, composite99DocValuesConsumer);
+            starTreesBuilder.buildDuringMerge(
+                metaOut,
+                dataOut,
+                starTreeSubsPerField,
+                composite99DocValuesConsumer,
+                mergeState,
+                fieldDocIdSetIteratorMap
+            );
         }
     }
 

diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/DateDimension.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/DateDimension.java
@@ -8,6 +8,7 @@
 
 package org.opensearch.index.compositeindex.datacube;
 
+import org.apache.lucene.index.DocValuesType;
 import org.opensearch.common.Rounding;
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.common.time.DateUtils;
@@ -98,6 +99,11 @@ public List<String> getSubDimensionNames() {
         return fields;
     }
 
+    @Override
+    public DocValuesType getDocValuesType() {
+        return DocValuesType.SORTED_NUMERIC;
+    }
+
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject("date_dimension");

diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/Dimension.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/Dimension.java
@@ -8,6 +8,7 @@
 
 package org.opensearch.index.compositeindex.datacube;
 
+import org.apache.lucene.index.DocValuesType;
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.core.xcontent.ToXContent;
 
@@ -42,4 +43,6 @@ public interface Dimension extends ToXContent {
      * Returns the list of dimension fields that represent the dimension
      */
     List<String> getSubDimensionNames();
+
+    DocValuesType getDocValuesType();
 }
diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/KeywordDimension.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/KeywordDimension.java
@@ -8,13 +8,15 @@
 
 package org.opensearch.index.compositeindex.datacube;
 
+import org.apache.lucene.index.DocValuesType;
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.index.mapper.CompositeDataCubeFieldType;
 
 import java.io.IOException;
 import java.util.List;
 import java.util.Objects;
+import java.util.function.Consumer;
 
 /**
  * Composite index keyword dimension class
@@ -41,16 +43,20 @@ public int getNumSubDimensions() {
     }
 
     @Override
-    public int setDimensionValues(Long value, Long[] dims, int index) {
-        dims[index++] = value;
-        return index;
+    public void setDimensionValues(Long value, Consumer<Long> dimSetter) {
+        dimSetter.accept(value);
     }
 
     @Override
-    public List<String> getDimensionFieldsNames() {
+    public List<String> getSubDimensionNames() {
         return List.of(field);
     }
 
+    @Override
+    public DocValuesType getDocValuesType() {
+        return DocValuesType.SORTED_SET;
+    }
+
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject();

diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/NumericDimension.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/NumericDimension.java
@@ -8,6 +8,7 @@
 
 package org.opensearch.index.compositeindex.datacube;
 
+import org.apache.lucene.index.DocValuesType;
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.index.mapper.CompositeDataCubeFieldType;
@@ -50,6 +51,11 @@ public List<String> getSubDimensionNames() {
         return List.of(field);
     }
 
+    @Override
+    public DocValuesType getDocValuesType() {
+        return DocValuesType.SORTED_NUMERIC;
+    }
+
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject();

diff --git a/server/src/main/java/org/opensearch/index/compositeindex/datacube/ReadDimension.java b/server/src/main/java/org/opensearch/index/compositeindex/datacube/ReadDimension.java
@@ -8,6 +8,7 @@
 
 package org.opensearch.index.compositeindex.datacube;
 
+import org.apache.lucene.index.DocValuesType;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.index.mapper.CompositeDataCubeFieldType;
 
@@ -24,9 +25,16 @@
 public class ReadDimension implements Dimension {
     public static final String READ = "read";
     private final String field;
+    private final DocValuesType docValuesType;
 
     public ReadDimension(String field) {
         this.field = field;
+        this.docValuesType = DocValuesType.SORTED_NUMERIC;
+    }
+
+    public ReadDimension(String field, DocValuesType docValuesType) {
+        this.field = field;
+        this.docValuesType = docValuesType;
     }
 
     public String getField() {
@@ -48,6 +56,11 @@ public List<String> getSubDimensionNames() {
         return List.of(field);
     }
 
+    @Override
+    public DocValuesType getDocValuesType() {
+        return docValuesType;
+    }
+
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject();
@@ -69,4 +82,5 @@ public boolean equals(Object o) {
     public int hashCode() {
         return Objects.hash(field);
     }
+
 }
diff --git a/...er/src/main/java/org/opensearch/index/compositeindex/datacube/startree/StarTreeField.java b/...er/src/main/java/org/opensearch/index/compositeindex/datacube/startree/StarTreeField.java
@@ -8,6 +8,7 @@
 
 package org.opensearch.index.compositeindex.datacube.startree;
 
+import org.apache.lucene.index.DocValuesType;
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.core.xcontent.ToXContent;
 import org.opensearch.core.xcontent.XContentBuilder;
@@ -33,6 +34,7 @@ public class StarTreeField implements ToXContent {
     private final List<Metric> metrics;
     private final StarTreeFieldConfiguration starTreeConfig;
     private final List<String> dimensionNames;
+    private final List<DocValuesType> dimensionDocValueTypes;
     private final List<String> metricNames;
 
     public StarTreeField(String name, List<Dimension> dimensions, List<Metric> metrics, StarTreeFieldConfiguration starTreeConfig) {
@@ -41,8 +43,12 @@ public StarTreeField(String name, List<Dimension> dimensions, List<Metric> metri
         this.metrics = metrics;
         this.starTreeConfig = starTreeConfig;
         dimensionNames = new ArrayList<>();
+        dimensionDocValueTypes = new ArrayList<>();
         for (Dimension dimension : dimensions) {
-            dimensionNames.addAll(dimension.getSubDimensionNames());
+            for (String dimensionName : dimension.getSubDimensionNames()) {
+                dimensionNames.add(dimensionName);
+                dimensionDocValueTypes.add(dimension.getDocValuesType());
+            }
         }
         metricNames = new ArrayList<>();
         for (Metric metric : metrics) {
@@ -64,6 +70,10 @@ public List<String> getDimensionNames() {
         return dimensionNames;
     }
 
+    public List<DocValuesType> getDimensionDocValueTypes() {
+        return dimensionDocValueTypes;
+    }
+
     public List<String> getMetricNames() {
         return metricNames;
     }