CNDB-14159: Fix DOC_LENGTHS offset calculations (#1745)

### What is the issue
Addresses some of the issues discovered by
riptano/cndb#14159

### What does this PR fix and why was it fixed

First I added tests that show the issue. Then, I fixed it by setting
`EC` will to still have the "bug" where the ComponentMetadata will point
to just after the header while ED and beyond will point to the header.

There are no places we use the header other than the SystemView to get
info about the length (which was wrong due to not including the header).
However, these do not affect the validation of the header, so I propose
that we leave them as is and just fix it header other than the
SystemView to get info about the length (which was wrong due to not
including the header). However, these do not affect the validation of
the header, so I propose that we leave them as is and just fix it in ED.

Author: michaeljmarshall

src/java/org/apache/cassandra/index/sai/disk/v1/DocLengthsReader.java

@@ -22,30 +22,38 @@
 
 import javax.annotation.concurrent.NotThreadSafe;
 
+import org.apache.cassandra.index.sai.disk.format.Version;
 import org.apache.cassandra.index.sai.disk.io.IndexInputReader;
 import org.apache.cassandra.index.sai.utils.IndexFileUtils;
 import org.apache.cassandra.index.sai.utils.SAICodecUtils;
 import org.apache.cassandra.io.util.FileHandle;
 import org.apache.cassandra.io.util.FileUtils;
-import org.apache.lucene.codecs.CodecUtil;
 
+/**
+ * Reads the component written by {@link org.apache.cassandra.index.sai.disk.v1.trie.DocLengthsWriter}.
+ */
 @NotThreadSafe
 public class DocLengthsReader implements Closeable
 {
     private final IndexInputReader input;
-    private final SegmentMetadata.ComponentMetadata componentMetadata;
+    private final long offset;
+    private final long upperBound;
 
-    public DocLengthsReader(FileHandle fileHandle, SegmentMetadata.ComponentMetadata componentMetadata)
+    public DocLengthsReader(FileHandle fileHandle, SegmentMetadata.ComponentMetadata componentMetadata, Version version)
     {
         this.input = IndexFileUtils.instance.openInput(fileHandle);
-        this.componentMetadata = componentMetadata;
+        // Version EC skipped the header in the doc lengths component metadata.
+        int headerAdjustment = Version.EC.equals(version) ? 0 : SAICodecUtils.headerSize();
+        this.offset = componentMetadata.offset + headerAdjustment;
+        // The offset + length get you the end of the file for all relevant versions.
+        this.upperBound = componentMetadata.offset + componentMetadata.length;
     }
 
     public int get(int rowID) throws IOException
     {
         // Account for header size in offset calculation
-        long position = componentMetadata.offset + (long) rowID * Integer.BYTES;
-        if (position >= componentMetadata.offset + componentMetadata.length)
+        long position = offset + (long) rowID * Integer.BYTES;
+        if (position >= upperBound)
             return 0;
         input.seek(position);
         return input.readInt();

src/java/org/apache/cassandra/index/sai/disk/v1/InvertedIndexSearcher.java

@@ -209,7 +209,7 @@ public CloseableIterator<PrimaryKeyWithSortKey> orderBy(Orderer orderer, Express
 
         var pkm = primaryKeyMapFactory.newPerSSTablePrimaryKeyMap();
         var merged = IntersectingPostingList.intersect(postingLists);
-        var docLengthsReader = new DocLengthsReader(docLengths, docLengthsMeta);
+        var docLengthsReader = new DocLengthsReader(docLengths, docLengthsMeta, version);
 
         // Wrap the iterator with resource management
         var it = new AbstractIterator<BM25Utils.DocTF>() { // Anonymous class extends AbstractIterator

src/java/org/apache/cassandra/index/sai/disk/v1/trie/DocLengthsWriter.java

@@ -23,6 +23,7 @@
 import org.agrona.collections.Int2IntHashMap;
 import org.apache.cassandra.index.sai.disk.format.IndexComponents;
 import org.apache.cassandra.index.sai.disk.format.IndexComponentType;
+import org.apache.cassandra.index.sai.disk.format.Version;
 import org.apache.cassandra.index.sai.disk.io.IndexOutputWriter;
 import org.apache.cassandra.index.sai.utils.SAICodecUtils;
 
@@ -33,10 +34,23 @@ public class DocLengthsWriter implements Closeable
 {
     private final IndexOutputWriter output;
 
+    private final long startOffset;
+
     public DocLengthsWriter(IndexComponents.ForWrite components) throws IOException
     {
         this.output = components.addOrGet(IndexComponentType.DOC_LENGTHS).openOutput(true);
-        SAICodecUtils.writeHeader(output);
+
+        // Version EC skipped the header in the doc lengths component metadata.
+        if (Version.EC.equals(components.version()))
+        {
+            SAICodecUtils.writeHeader(output);
+            startOffset = output.getFilePointer();
+        }
+        else
+        {
+            startOffset = output.getFilePointer();
+            SAICodecUtils.writeHeader(output);
+        }
     }
 
     public void writeDocLengths(Int2IntHashMap lengths) throws IOException
@@ -65,6 +79,14 @@ public long getFilePointer()
         return output.getFilePointer();
     }
 
+    /**
+     * @return file pointer where index structure begins (before header)
+     */
+    public long getStartOffset()
+    {
+        return startOffset;
+    }
+
     @Override
     public void close() throws IOException
     {

src/java/org/apache/cassandra/index/sai/disk/v1/trie/InvertedIndexWriter.java

@@ -103,7 +103,7 @@ public SegmentMetadata.ComponentMetadataMap writeAll(TermsIterator terms, Int2In
         // Write doc lengths
         if (docLengthsWriter != null)
         {
-            long docLengthsOffset = docLengthsWriter.getFilePointer();
+            long docLengthsOffset = docLengthsWriter.getStartOffset();
             docLengthsWriter.writeDocLengths(docLengths);
             long docLengthsLength = docLengthsWriter.getFilePointer() - docLengthsOffset;
             components.put(IndexComponentType.DOC_LENGTHS, -1, docLengthsOffset, docLengthsLength);

test/unit/org/apache/cassandra/index/sai/SAITester.java

@@ -73,6 +73,7 @@
 import org.apache.cassandra.index.sai.disk.format.IndexComponentType;
 import org.apache.cassandra.index.sai.disk.format.IndexDescriptor;
 import org.apache.cassandra.index.sai.disk.format.Version;
+import org.apache.cassandra.index.sai.disk.v1.SegmentBuilder;
 import org.apache.cassandra.index.sai.plan.QueryController;
 import org.apache.cassandra.index.sai.plan.StorageAttachedIndexQueryPlan;
 import org.apache.cassandra.index.sai.plan.StorageAttachedIndexSearcher;
@@ -225,6 +226,14 @@ public void resetQueryOptimizationLevel() throws Throwable
     {
         // Enable the optimizer by default. If there are any tests that need to disable it, they can do so explicitly.
         QueryController.QUERY_OPT_LEVEL = 1;
+
+    }
+
+    @Before
+    public void resetLastValidSegmentRowId() throws Throwable
+    {
+        // Don't want this setting to impact peer tests
+        SegmentBuilder.updateLastValidSegmentRowId(-1);
     }
 
     @After

test/unit/org/apache/cassandra/index/sai/virtual/SegmentsSystemViewTest.java

@@ -22,8 +22,15 @@
 import java.util.Map;
 
 import com.google.common.collect.ImmutableList;
+import org.junit.Before;
 import org.junit.BeforeClass;
 import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import java.util.Collection;
+import java.util.stream.Collectors;
+import org.apache.cassandra.index.sai.disk.format.Version;
+import org.apache.cassandra.index.sai.SAIUtil;
 
 import org.apache.cassandra.cql3.UntypedResultSet;
 import org.apache.cassandra.db.ColumnFamilyStore;
@@ -40,6 +47,7 @@
 import org.apache.cassandra.index.sai.disk.format.IndexDescriptor;
 import org.apache.cassandra.index.sai.disk.v1.SegmentBuilder;
 import org.apache.cassandra.index.sai.disk.v1.SegmentMetadata;
+import org.apache.cassandra.index.sai.utils.SAICodecUtils;
 import org.apache.cassandra.index.sai.utils.TypeUtil;
 import org.apache.cassandra.io.sstable.format.SSTableReader;
 import org.apache.cassandra.schema.SchemaConstants;
@@ -52,8 +60,24 @@
 /**
  * Tests the virtual table exposing SSTable index segment metadata.
  */
+@RunWith(Parameterized.class)
 public class SegmentsSystemViewTest extends SAITester
 {
+    @Parameterized.Parameter
+    public Version version;
+
+    @Parameterized.Parameters(name = "{0}")
+    public static Collection<Object[]> data()
+    {
+        return Version.ALL.stream().map(v -> new Object[]{ v }).collect(Collectors.toList());
+    }
+
+    @Before
+    public void setupVersion()
+    {
+        SAIUtil.setCurrentVersion(version);
+    }
+
     private static final String SELECT = String.format("SELECT %s, %s, %s, %s " +
                                                        "FROM %s.%s WHERE %s = '%s' AND %s = ?",
                                                        SegmentsSystemView.SEGMENT_ROW_ID_OFFSET,
@@ -210,6 +234,17 @@ private HashMap<String, Long> indexFileLengths(String table) throws Exception
                 {
                     addComponentSizeToMap(lengths, IndexComponentType.TERMS_DATA, index.getIndexContext(), indexDescriptor);
                     addComponentSizeToMap(lengths, IndexComponentType.POSTING_LISTS, index.getIndexContext(), indexDescriptor);
+                    if (version.onOrAfter(Version.BM25_EARLIEST))
+                    {
+                        addComponentSizeToMap(lengths, IndexComponentType.DOC_LENGTHS, index.getIndexContext(), indexDescriptor);
+                        // Version EC does not count the length of the segment header in the DOC_LENGTHS file, so
+                        // we do a special adjustment here
+                        if (version.equals(Version.EC))
+                        {
+                            var error = sstableIndex.getSegments().size() * SAICodecUtils.headerSize();
+                            lengths.computeIfPresent(IndexComponentType.DOC_LENGTHS.name(), (typeName, acc) -> acc - error);
+                        }
+                    }
                 }
                 else
                 {