PQVectors::encodeAndBuild with ordinals mapping is too restrictive

[reta]

Introduction

The PQVectors::encodeAndBuild with ordinals mapping provides the way to map vectors [0 ... vectorCount) from RandomAccessVectorValues using custom ordinalsMapping function. It works great in general but have a hard restriction that ordinalsMapping function cannot return any value greater that vectorCount.

Problem statement

The fact that ordinalsMapping function cannot return any value greater that vectorCount is too restrictive for a few reasons:

• the PQ vector ordinal space is created as contiguous [0 ... vectorCount) range (no ordinalsMapping involvement)
• the value of ordinalsMapping function is not relevant to PQ vector space
• the value of ordinalsMapping is only relevant to RandomAccessVectorValues instance (ravv)

Here is the proof from current implementation (this is a bit simplified version to get rid of thread locals etc):

IntStream.range(0, vectorCount)
      .parallel()
      .forEach(ordinal -> {
          // Retrieve the slice and mutate it.
          var slice = PQVectors.get(chunks, ordinal, layout.fullChunkVectors, pq.getSubspaceCount());
          var vector = ravv.getVector(ordinalsMapping.applyAsInt(ordinal));
          if (vector != null)
              pq.encodeTo(vector, slice);
          else
              slice.zero();
      }))

The key statement here is var vector = ravv.getVector(ordinalsMapping.applyAsInt(ordinal));: there is no direct / indirect dependency between vectorCount and ravv size here, the only things that matters is that value returned by ordinalsMapping.applyAsInt(ordinal) is a valid input to ravv.

Why it matters

In opensearch-jvector plugin we have so called leading segment merge, where the simple idea is employed:

• use largest segment graph as is (live + deleted vectors)
• add live vectors from all other segments into the graph
• clean (compact) and flush to disk as a new segment

In case of PQ vectors, we are doing the following:

• use PQ vectors from leading segment + all other segments
• after compaction, remap PQ vectors by filtering out deleted vectors

The latter leads to the issue:

• we use ravv from previous step (which has live + deleted vectors from leading segment, live vectors from other segments)
• the vectorCount is mostly guaranteed to be < ravv.size()
• the ordinalsMapping only maps the live segments (no deleted vectors), and is mostly guaranteed to return ordinal values > vectorCount
• PQVectors::encodeAndBuild fails with validation exception prematurely

Solution space

• remove the ordinalsMapping function pre-check against vectorCount (could potentially lead to regressions)

• introduce ordinalsMapping with validation support (this is non-breaking change and could preserve existing validation logic untouched)

  interface ValidatableOrdinalMapper extends IntUnaryOperator {
      void validate(int minOrd, int maxOrd);
  }

  Used liked this:

   if (ordinalsMapping instanceof ValidatableOrdinalMapper v) {
       validate.validate(0, vectorCount);  
   } else {
       // existing check
   }
  

• do nothing (we could replicate PQVectors::encodeAndBuild internally in the plugin)

[reta]

CC @MarkWolters