addressing comments

Signed-off-by: Bharathwaj G <[email protected]>

server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/ByteListBackedBitset.java → server/src/main/java/org/opensearch/common/util/ByteArrayBackedBitset.java

@@ -6,30 +6,31 @@
  * compatible open source license.
  */
 
-package org.opensearch.index.compositeindex.datacube.startree.utils;
+package org.opensearch.common.util;
 
+import org.apache.lucene.store.IndexInput;
 import org.apache.lucene.store.IndexOutput;
 import org.apache.lucene.store.RandomAccessInput;
 
 import java.io.IOException;
 
 /**
- * This is used to store and retrieve null values in the bitset.
+ * A bitset backed by a byte array. This will initialize and set bits in the byte array based on the index.
  */
-public class ByteListBackedBitset {
+public class ByteArrayBackedBitset {
     private final byte[] byteArray;
 
     /**
      * Constructor which uses an on heap list. This should be using during construction of the bitset.
      */
-    public ByteListBackedBitset(int capacity) {
+    public ByteArrayBackedBitset(int capacity) {
         byteArray = new byte[capacity];
     }
 
     /**
-     * Constructor which set the Lucene's IndexInput to read the bitset into a read-only buffer.
+     * Constructor which set the Lucene's RandomAccessInput to read the bitset into a read-only buffer.
      */
-    public ByteListBackedBitset(RandomAccessInput in, long offset, int length) throws IOException {
+    public ByteArrayBackedBitset(RandomAccessInput in, long offset, int length) throws IOException {
         byteArray = new byte[length];
         int i = 0;
         while (i < length) {
@@ -38,6 +39,23 @@ public ByteListBackedBitset(RandomAccessInput in, long offset, int length) throw
         }
     }
 
+    /**
+     * Constructor which set the Lucene's IndexInput to read the bitset into a read-only buffer.
+     */
+    public ByteArrayBackedBitset(IndexInput in, int length) throws IOException {
+        byteArray = new byte[length];
+        int i = 0;
+        while (i < length) {
+            byteArray[i] = in.readByte();
+            i++;
+        }
+    }
+
+    /**
+     * Sets the bit at the given index to 1.
+     * Each byte can indicate 8 bits, so the index is divided by 8 to get the byte array index.
+     * @param index the index to set the bit
+     */
     public void set(int index) {
         int byteArrIndex = index >> 3;
         byteArray[byteArrIndex] |= (byte) (1 << (index & 7));
@@ -52,6 +70,11 @@ public int write(IndexOutput output) throws IOException {
         return numBytes;
     }
 
+    /**
+     * Retrieves whether the bit is set or not at the given index.
+     * @param index the index to look up for the bit
+     * @return true if bit is set, false otherwise
+     */
     public boolean get(int index) throws IOException {
         int byteArrIndex = index >> 3;
         return (byteArray[byteArrIndex] & (1 << (index & 7))) != 0;

server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/AbstractDocumentsFileManager.java

@@ -118,19 +118,18 @@ protected int writeMetrics(StarTreeDocument starTreeDocument, IndexOutput output
      *
      * @param input   RandomAccessInput
      * @param offset  Offset in the file
-     * @param shouldReadAggregatedDocs boolean to indicate if aggregated star tree docs should be read
+     * @param isAggregatedDoc boolean to indicate if aggregated star tree docs should be read
      * @return StarTreeDocument
      * @throws IOException IOException in case of I/O errors
      */
-    protected StarTreeDocument readStarTreeDocument(RandomAccessInput input, long offset, boolean shouldReadAggregatedDocs)
-        throws IOException {
+    protected StarTreeDocument readStarTreeDocument(RandomAccessInput input, long offset, boolean isAggregatedDoc) throws IOException {
         int dimSize = starTreeField.getDimensionsOrder().size();
         Long[] dimensions = new Long[dimSize];
         long initialOffset = offset;
         offset = readDimensions(dimensions, input, offset);
 
         Object[] metrics = new Object[numMetrics];
-        offset = readMetrics(input, offset, numMetrics, metrics, shouldReadAggregatedDocs);
+        offset = readMetrics(input, offset, numMetrics, metrics, isAggregatedDoc);
         assert (offset - initialOffset) == docSizeInBytes;
         return new StarTreeDocument(dimensions, metrics);
     }
@@ -155,7 +154,7 @@ protected long readDimensions(Long[] dimensions, RandomAccessInput input, long o
     /**
      * Read star tree metrics from file
      */
-    protected long readMetrics(RandomAccessInput input, long offset, int numMetrics, Object[] metrics, boolean shouldReadAggregatedDocs)
+    protected long readMetrics(RandomAccessInput input, long offset, int numMetrics, Object[] metrics, boolean isAggregatedDoc)
         throws IOException {
         for (int i = 0; i < numMetrics; i++) {
             switch (metricAggregatorInfos.get(i).getValueAggregators().getAggregatedValueType()) {
@@ -165,7 +164,7 @@ protected long readMetrics(RandomAccessInput input, long offset, int numMetrics,
                     break;
                 case DOUBLE:
                     long val = input.readLong(offset);
-                    if (shouldReadAggregatedDocs) {
+                    if (isAggregatedDoc) {
                         metrics[i] = StarTreeNumericTypeConverters.sortableLongtoDouble(val);
                     } else {
                         metrics[i] = val;
@@ -193,7 +192,7 @@ protected long readMetrics(RandomAccessInput input, long offset, int numMetrics,
     /**
      * Read star tree document from file based on doc id
      */
-    public abstract StarTreeDocument readStarTreeDocument(int docId, boolean isMerge) throws IOException;
+    public abstract StarTreeDocument readStarTreeDocument(int docId, boolean isAggregatedDoc) throws IOException;
 
     /**
      * Read star document dimensions from file based on doc id

server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/SegmentDocsFileManager.java

@@ -70,9 +70,9 @@ private void maybeInitializeSegmentInput() throws IOException {
     }
 
     @Override
-    public StarTreeDocument readStarTreeDocument(int docId, boolean isMerge) throws IOException {
+    public StarTreeDocument readStarTreeDocument(int docId, boolean isAggregatedDoc) throws IOException {
         maybeInitializeSegmentInput();
-        return readStarTreeDocument(segmentRandomInput, (long) docId * docSizeInBytes, isMerge);
+        return readStarTreeDocument(segmentRandomInput, (long) docId * docSizeInBytes, isAggregatedDoc);
     }
 
     @Override

server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/builder/StarTreeDocsFileManager.java

@@ -199,13 +199,13 @@ private void loadStarTreeDocumentFile(int docId) throws IOException {
      *    If the operation is only for reading existing documents, a new file is not created.
      */
     private void closeAndMaybeCreateNewFile(boolean shouldCreateFileForAppend, int numStarTreeDocs) throws IOException {
-        if (starTreeDocsFileOutput != null) {
-            IOUtils.close(starTreeDocsFileOutput);
-        }
         currBytes = 0;
         if (starTreeDocsFileOutput != null) {
             fileToEndDocIdMap.put(starTreeDocsFileOutput.getName(), numStarTreeDocs);
         }
+        if (starTreeDocsFileOutput != null) {
+            IOUtils.close(starTreeDocsFileOutput);
+        }
         if (shouldCreateFileForAppend) {
             starTreeDocsFileOutput = createStarTreeDocumentsFileOutput();
             if (fileToEndDocIdMap.size() >= fileCountMergeThreshold) {

server/src/main/java/org/opensearch/index/compositeindex/datacube/startree/utils/StarTreeDocumentBitSetUtil.java

@@ -10,6 +10,7 @@
 
 import org.apache.lucene.store.IndexOutput;
 import org.apache.lucene.store.RandomAccessInput;
+import org.opensearch.common.util.ByteArrayBackedBitset;
 
 import java.io.IOException;
 import java.util.function.Function;
@@ -27,7 +28,7 @@ public class StarTreeDocumentBitSetUtil {
      * @throws IOException if an I/O error occurs while writing to the output stream
      */
     public static int writeBitSet(Object[] array, IndexOutput output) throws IOException {
-        ByteListBackedBitset bitset = new ByteListBackedBitset(getLength(array));
+        ByteArrayBackedBitset bitset = new ByteArrayBackedBitset(getLength(array));
         for (int i = 0; i < array.length; i++) {
             if (array[i] == null) {
                 bitset.set(i);
@@ -41,7 +42,7 @@ public static int writeBitSet(Object[] array, IndexOutput output) throws IOExcep
      */
     public static int readBitSet(RandomAccessInput input, long offset, Object[] array, Function<Integer, Object> identityValueSupplier)
         throws IOException {
-        ByteListBackedBitset bitset = new ByteListBackedBitset(input, offset, getLength(array));
+        ByteArrayBackedBitset bitset = new ByteArrayBackedBitset(input, offset, getLength(array));
         for (int i = 0; i < array.length; i++) {
             if (bitset.get(i)) {
                 array[i] = identityValueSupplier.apply(i);