Star tree iterator changes

Signed-off-by: Bharathwaj G <bharath78910@gmail.com>
bharath-techie · Sep 17, 2024 · ade1536 · ade1536
1 parent e3bbc74
commit ade1536
Show file tree

Hide file tree

Showing 7 changed files with 142 additions and 75 deletions.
diff --git a/...va/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilder.java b/...va/org/opensearch/index/compositeindex/datacube/startree/builder/BaseStarTreeBuilder.java
@@ -34,6 +34,7 @@
 import org.opensearch.index.compositeindex.datacube.startree.node.InMemoryTreeNode;
 import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNodeType;
 import org.opensearch.index.compositeindex.datacube.startree.utils.SequentialDocValuesIterator;
+import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedNumericStarTreeValuesIterator;
 import org.opensearch.index.mapper.DocCountFieldMapper;
 import org.opensearch.index.mapper.FieldMapper;
 import org.opensearch.index.mapper.FieldValueConverter;
@@ -193,7 +194,9 @@ public List<SequentialDocValuesIterator> getMetricReaders(SegmentWriteState stat
                         metricFieldInfo = getFieldInfo(metric.getField(), DocValuesType.SORTED_NUMERIC);
                     }
                     metricReader = new SequentialDocValuesIterator(
-                        fieldProducerMap.get(metricFieldInfo.name).getSortedNumeric(metricFieldInfo)
+                        new SortedNumericStarTreeValuesIterator(
+                            fieldProducerMap.get(metricFieldInfo.name).getSortedNumeric(metricFieldInfo)
+                        )
                     );
                 }
                 metricReaders.add(metricReader);
@@ -228,7 +231,7 @@ public void build(
                 dimensionFieldInfo = getFieldInfo(dimension, DocValuesType.SORTED_NUMERIC);
             }
             dimensionReaders[i] = new SequentialDocValuesIterator(
-                fieldProducerMap.get(dimensionFieldInfo.name).getSortedNumeric(dimensionFieldInfo)
+                new SortedNumericStarTreeValuesIterator(fieldProducerMap.get(dimensionFieldInfo.name).getSortedNumeric(dimensionFieldInfo))
             );
         }
         Iterator<StarTreeDocument> starTreeDocumentIterator = sortAndAggregateSegmentDocuments(dimensionReaders, metricReaders);
@@ -672,7 +675,7 @@ private SequentialDocValuesIterator getIteratorForNumericField(
         SequentialDocValuesIterator sequentialDocValuesIterator;
         assert fieldProducerMap.containsKey(fieldInfo.name);
         sequentialDocValuesIterator = new SequentialDocValuesIterator(
-            DocValues.singleton(fieldProducerMap.get(fieldInfo.name).getNumeric(fieldInfo))
+            new SortedNumericStarTreeValuesIterator(DocValues.singleton(fieldProducerMap.get(fieldInfo.name).getNumeric(fieldInfo)))
         );
         return sequentialDocValuesIterator;
     }

diff --git a/...main/java/org/opensearch/index/compositeindex/datacube/startree/index/StarTreeValues.java b/...main/java/org/opensearch/index/compositeindex/datacube/startree/index/StarTreeValues.java
@@ -12,7 +12,6 @@
 import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.SortedNumericDocValues;
-import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.store.IndexInput;
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.index.compositeindex.CompositeIndexMetadata;
@@ -25,6 +24,8 @@
 import org.opensearch.index.compositeindex.datacube.startree.fileformats.meta.StarTreeMetadata;
 import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeFactory;
 import org.opensearch.index.compositeindex.datacube.startree.node.StarTreeNode;
+import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedNumericStarTreeValuesIterator;
+import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.StarTreeValuesIterator;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -61,12 +62,12 @@ public class StarTreeValues implements CompositeIndexValues {
     /**
      * A map containing suppliers for DocIdSetIterators for dimensions.
      */
-    private final Map<String, Supplier<DocIdSetIterator>> dimensionDocValuesIteratorMap;
+    private final Map<String, Supplier<StarTreeValuesIterator>> dimensionDocValuesIteratorMap;
 
     /**
      * A map containing suppliers for DocIdSetIterators for metrics.
      */
-    private final Map<String, Supplier<DocIdSetIterator>> metricDocValuesIteratorMap;
+    private final Map<String, Supplier<StarTreeValuesIterator>> metricDocValuesIteratorMap;
 
     /**
      * A map containing attributes associated with the star tree values.
@@ -91,8 +92,8 @@ public class StarTreeValues implements CompositeIndexValues {
     public StarTreeValues(
         StarTreeField starTreeField,
         StarTreeNode root,
-        Map<String, Supplier<DocIdSetIterator>> dimensionDocValuesIteratorMap,
-        Map<String, Supplier<DocIdSetIterator>> metricDocValuesIteratorMap,
+        Map<String, Supplier<StarTreeValuesIterator>> dimensionDocValuesIteratorMap,
+        Map<String, Supplier<StarTreeValuesIterator>> metricDocValuesIteratorMap,
         Map<String, String> attributes,
         StarTreeMetadata compositeIndexMetadata
     ) {
@@ -162,7 +163,7 @@ public StarTreeValues(
                             dimensionSortedNumericDocValues = compositeDocValuesProducer.getSortedNumeric(dimensionfieldInfo);
                         }
                     }
-                    return getSortedNumericDocValues(dimensionSortedNumericDocValues);
+                    return new SortedNumericStarTreeValuesIterator(getSortedNumericDocValues(dimensionSortedNumericDocValues));
                 } catch (IOException e) {
                     throw new RuntimeException("Error loading dimension DocIdSetIterator", e);
                 }
@@ -186,7 +187,7 @@ public StarTreeValues(
                                 metricSortedNumericDocValues = compositeDocValuesProducer.getSortedNumeric(metricFieldInfo);
                             }
                         }
-                        return getSortedNumericDocValues(metricSortedNumericDocValues);
+                        return new SortedNumericStarTreeValuesIterator(getSortedNumericDocValues(metricSortedNumericDocValues));
                     } catch (IOException e) {
                         throw new RuntimeException("Error loading metric DocIdSetIterator", e);
                     }
@@ -244,7 +245,7 @@ public Map<String, String> getAttributes() {
      * @param dimension The name of the dimension.
      * @return The DocIdSetIterator for the specified dimension.
      */
-    public DocIdSetIterator getDimensionDocIdSetIterator(String dimension) {
+    public StarTreeValuesIterator getDimensionDocIdSetIterator(String dimension) {
 
         if (dimensionDocValuesIteratorMap.containsKey(dimension)) {
             return dimensionDocValuesIteratorMap.get(dimension).get();
@@ -259,7 +260,7 @@ public DocIdSetIterator getDimensionDocIdSetIterator(String dimension) {
      * @param fullyQualifiedMetricName The fully qualified name of the metric.
      * @return The DocIdSetIterator for the specified fully qualified metric name.
      */
-    public DocIdSetIterator getMetricDocIdSetIterator(String fullyQualifiedMetricName) {
+    public StarTreeValuesIterator getMetricDocIdSetIterator(String fullyQualifiedMetricName) {
 
         if (metricDocValuesIteratorMap.containsKey(fullyQualifiedMetricName)) {
             return metricDocValuesIteratorMap.get(fullyQualifiedMetricName).get();

diff --git a/.../opensearch/index/compositeindex/datacube/startree/utils/SequentialDocValuesIterator.java b/.../opensearch/index/compositeindex/datacube/startree/utils/SequentialDocValuesIterator.java
@@ -12,12 +12,15 @@
 import org.apache.lucene.index.SortedNumericDocValues;
 import org.apache.lucene.search.DocIdSetIterator;
 import org.opensearch.common.annotation.ExperimentalApi;
+import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.SortedNumericStarTreeValuesIterator;
+import org.opensearch.index.compositeindex.datacube.startree.utils.iterator.StarTreeValuesIterator;
 
 import java.io.IOException;
 
 /**
  * Coordinates the reading of documents across multiple DocIdSetIterators.
  * It encapsulates a single DocIdSetIterator and maintains the latest document ID and its associated value.
+ *
  * @opensearch.experimental
  */
 @ExperimentalApi
@@ -26,12 +29,7 @@ public class SequentialDocValuesIterator {
     /**
      * The doc id set iterator associated for each field.
      */
-    private final DocIdSetIterator docIdSetIterator;
-
-    /**
-     * The value associated with the latest document.
-     */
-    private Long docValue;
+    private final StarTreeValuesIterator starTreeValuesIterator;
 
     /**
      * The id of the latest document.
@@ -44,7 +42,15 @@ public class SequentialDocValuesIterator {
      * @param docIdSetIterator the DocIdSetIterator to be associated with this instance
      */
     public SequentialDocValuesIterator(DocIdSetIterator docIdSetIterator) {
-        this.docIdSetIterator = docIdSetIterator;
+        if (docIdSetIterator instanceof SortedNumericDocValues) {
+            this.starTreeValuesIterator = new SortedNumericStarTreeValuesIterator(docIdSetIterator);
+        } else {
+            throw new IllegalStateException("Unsupported Iterator requested for SequentialDocValuesIterator");
+        }
+    }
+
+    public SequentialDocValuesIterator(StarTreeValuesIterator starTreeValuesIterator) {
+        this.starTreeValuesIterator = starTreeValuesIterator;
     }
 
     /**
@@ -65,37 +71,27 @@ private void setDocId(int docId) {
         this.docId = docId;
     }
 
-    /**
-     * Returns the DocIdSetIterator associated with this instance.
-     *
-     * @return the DocIdSetIterator associated with this instance
-     */
-    public DocIdSetIterator getDocIdSetIterator() {
-        return docIdSetIterator;
-    }
-
     public int nextDoc(int currentDocId) throws IOException {
         // if doc id stored is less than or equal to the requested doc id , return the stored doc id
         if (docId >= currentDocId) {
             return docId;
         }
-        setDocId(this.docIdSetIterator.nextDoc());
+        setDocId(this.starTreeValuesIterator.nextEntry());
         return docId;
     }
 
     public Long value(int currentDocId) throws IOException {
-        if (this.getDocIdSetIterator() instanceof SortedNumericDocValues) {
-            SortedNumericDocValues sortedNumericDocValues = (SortedNumericDocValues) this.getDocIdSetIterator();
+        if (starTreeValuesIterator instanceof SortedNumericStarTreeValuesIterator) {
             if (currentDocId < 0) {
                 throw new IllegalStateException("invalid doc id to fetch the next value");
             }
             if (currentDocId == DocIdSetIterator.NO_MORE_DOCS) {
-                throw new IllegalStateException("DocValuesIterator is already exhausted");
+                throw new IllegalStateException("StarTreeValuesIterator is already exhausted");
             }
             if (docId == DocIdSetIterator.NO_MORE_DOCS || docId != currentDocId) {
                 return null;
             }
-            return sortedNumericDocValues.nextValue();
+            return ((SortedNumericStarTreeValuesIterator) starTreeValuesIterator).nextValue();
 
         } else {
             throw new IllegalStateException("Unsupported Iterator requested for SequentialDocValuesIterator");

diff --git a/.../compositeindex/datacube/startree/utils/iterator/SortedNumericStarTreeValuesIterator.java b/.../compositeindex/datacube/startree/utils/iterator/SortedNumericStarTreeValuesIterator.java
@@ -0,0 +1,27 @@
+/*
+ * 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.opensearch.index.compositeindex.datacube.startree.utils.iterator;
+
+import org.apache.lucene.index.SortedNumericDocValues;
+import org.apache.lucene.search.DocIdSetIterator;
+import org.opensearch.common.annotation.ExperimentalApi;
+
+import java.io.IOException;
+
+@ExperimentalApi
+public class SortedNumericStarTreeValuesIterator extends StarTreeValuesIterator {
+
+    public SortedNumericStarTreeValuesIterator(DocIdSetIterator docIdSetIterator) {
+        super(docIdSetIterator);
+    }
+
+    public long nextValue() throws IOException {
+        return ((SortedNumericDocValues) docIdSetIterator).nextValue();
+    }
+}
diff --git a/...nsearch/index/compositeindex/datacube/startree/utils/iterator/StarTreeValuesIterator.java b/...nsearch/index/compositeindex/datacube/startree/utils/iterator/StarTreeValuesIterator.java
@@ -0,0 +1,40 @@
+/*
+ * 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.opensearch.index.compositeindex.datacube.startree.utils.iterator;
+
+import org.apache.lucene.search.DocIdSetIterator;
+import org.opensearch.common.annotation.ExperimentalApi;
+
+import java.io.IOException;
+
+@ExperimentalApi
+public abstract class StarTreeValuesIterator {
+
+    protected final DocIdSetIterator docIdSetIterator;
+
+    public StarTreeValuesIterator(DocIdSetIterator docIdSetIterator) {
+        this.docIdSetIterator = docIdSetIterator;
+    }
+
+    public int entryId() {
+        return docIdSetIterator.docID();
+    }
+
+    public int nextEntry() throws IOException {
+        return docIdSetIterator.nextDoc();
+    }
+
+    public int advance(int target) throws IOException {
+        return docIdSetIterator.advance(target);
+    }
+
+    public long cost() {
+        return docIdSetIterator.cost();
+    }
+}