Move Points from a visitor API to a cursor-style API? [LUCENE-9619]

[asfimport]

Points' visitor API work well but there are a couple things we could make better if we moved to a cursor API, e.g.

• Term queries could return a DocIdSetIterator without having to materialize a BitSet.
• Nearest-neighbor search could work on top of the regular API instead of casting to BKDReader https://github.com/apache/lucene-solr/blob/6a7131ee246d700c2436a85ddc537575de2aeacf/lucene/sandbox/src/java/org/apache/lucene/sandbox/document/FloatPointNearestNeighbor.java#L296
• We could optimize counting the number of matches of a query by adding the number of points in a leaf without visiting documents where there are no deleted documents and a leaf fully matches the query.

---

Migrated from LUCENE-9619 by Adrien Grand (@jpountz), updated Dec 06 2021
Pull requests: #486

[asfimport]

Adrien Grand (@jpountz) (migrated from JIRA)

Here is what I have in mind in terms of API. The main downside compared to today is that it makes more assumptions about how points are implemented under the hood, e.g. it suggests a tree structure, which the current API doesn't. But I like that it would give us more control over how matching is performed as mentioned in the issue description. I still tried to push too many requirements to possible implementations, e.g. not enforcing an arity of 2 on inner nodes and not enforcing that the tree is balanced.

import java.io.IOException;
import java.util.function.IntConsumer;

import org.apache.lucene.search.DocIdSetIterator;

public abstract class PointValues {

  /* Global statistics, that don't change when moving from a node to another node. */

  /** Returns how many dimensions are represented in the values */
  public abstract int getNumDimensions() throws IOException;

  /** Returns how many dimensions are used for the index */
  public abstract int getNumIndexDimensions() throws IOException;

  /** Returns the number of bytes per dimension */
  public abstract int getBytesPerDimension() throws IOException;

  /** Return the total number of documents that have a value for this field, across all nodes. */
  public abstract int getTotalDocCount();

  /* Per-node statistics */
  
  /** Return the minimum packed value of the current node. */
  public abstract byte[] getMinPackedValue();

  /** Return the maximum packed value of the current node. */
  public abstract byte[] getMaxPackedValue();

  /** Return the total number of points under the current node. On the root node this returns the
   *  total number of points in the field on the current segment. */
  public abstract long size();

  /* API to walk the tree. */

  /** Move to the first child node and return {`@code` true} upon success. Returns {`@code` false} for
   *  leaf nodes and {`@code` true} otherwise. */
  public abstract boolean moveToChild();

  /** Move to the parent node and return {`@code` true} upon success. Returns {`@code` false} for the
   *  root node and {`@code` true} otherwise. */
  public abstract boolean moveToParent();

  /** Move to the next sibling node and return {`@code` true} upon success. Returns {`@code` false} if
   *  the current node has no more siblings. */
  public abstract boolean moveToSibling();

  /** A visitor for the content of the tree. */
  `@FunctionalInterface`
  public interface IntersectVisitor {

    /** Called for all documents in a leaf cell that crosses the query.  The consumer
     *  should scrutinize the packedValue to decide whether to accept it.  In the 1D case,
     *  values are visited in increasing order, and in the case of ties, in increasing
     *  docID order. */
    void visit(int docID, byte[] packedValue) throws IOException;

    /** Similar to {`@link` IntersectVisitor#visit(int, byte[])} but in this case the packedValue
     * can have more than one docID associated to it. The provided iterator should not escape the
     * scope of this method so that implementations of PointValues are free to reuse it,*/
    default void visit(DocIdSetIterator iterator, byte[] packedValue) throws IOException {
      int docID;
      while ((docID = iterator.nextDoc()) != DocIdSetIterator.NO_MORE_DOCS) {
        visit(docID, packedValue);
      }
    }
  }

  /** Visit all (document,value) pairs under the current node. {`@link` IntersectVisitor#visit} will
   *  be called {`@link` #size()} times. */
  public abstract void intersect(IntersectVisitor visitor);

  /** Visit all documents under the current node. {`@link` IntConsumer#accept} will be called
   *  {`@link` #size()} times. */
  public abstract void intersectAll(IntConsumer visitor);
}

Opinions?

[asfimport]

Ignacio Vera (@iverase) (migrated from JIRA)

Not sure if this API allows to be wrapped as a DocIdSetIterator. For example you cannot control how many documents you visit once you call IntersectsAll which can be potentially million of documents. This API seems to be cursor like only for inner nodes but not for leaf nodes, is it worthy the complexity of extending it so use can control how we visit documents on the leaf nodes? 

 

[asfimport]

Adrien Grand (@jpountz) (migrated from JIRA)

Thanks for the feedback @iverase!

I initially wanted to remove the visitor pattern entirely, but this made it challenging to retain some optimizations we have like the one we have to do only one comparison in case multiple documents share the same value: https://github.com/apache/lucene-solr/blob/af47cb7bcdd4eb10263a0586474c6e255307bbbb/lucene/core/src/java/org/apache/lucene/index/PointValues.java#L219-L224.

As far as implementing a DocIdSetIterator is concerned, my thinking was that this API could be use to fill an int[] buffer only one leaf at a time, so we wouldn't allocate more than an int[512] with the current Points file format. This wouldn't provide skipping capabilities, but at least we wouldn't need to maintain a giant int[] or BitSet.

[asfimport]

Ignacio Vera (@iverase) (migrated from JIRA)

my thinking was that this API could be use to fill an int[] buffer only one leaf at a time

 

With the visitor pattern there is no control to visit one leaf at a time. I am wondering if it make more sense to have and IntersectsAll method that just return a DocIdIdIterator. In addition the Intersects method can return a DocSetIterator as well and accept a Predicate function. Something like:

 

 /** Visit all (document,value) pairs under the current node and return only the ones that match the given predicate. The iterator will contain at most {`@link` #size()} elements. */
public abstract DocIdSetIterator intersect(Predicate<byte[]> matchesPredicate);

/** Return all documents under the current node. The iterator will contain {`@link` #size()} elements. */ 
public abstract DocIdSetIterator intersectAll();

 

I am actually more concern about the optimisation where we added a exact bounding box to the leaf nodes as it needs to be treated carefully.

[asfimport]

Adrien Grand (@jpountz) (migrated from JIRA)

The difference between what I'm proposing and the current API is that the intersect methods only look at data that is below the current node. So it should be possible to use the moveTo* methods to walk the tree in depth-first order and only call intersect* on the leaf nodes?

 My main concern about reusing DocIdSetIterator is that this requires sorting doc IDs since this si part of the contract of DocIdSetIterator. However I like your idea of using a predicate so maybe void intersect(IntersectVisitor visitor) could become void intersect(Predicate<byte[]> predicate, IntConsumer visitor)?

[asfimport]

Ignacio Vera (@iverase) (migrated from JIRA)

I see, yes DocIdSetIterator is not an option then.

so you manually visit each leaf node for IntersectsAll to make sure you only hold documents for one leaf at a time. The only frustrating thing of that approach is that you might need to copy the int array generated by the tree into another array in the IntConsumer? But anyway I see your point.

Yes, I think the predicate is a good option to avoid the intersectVisitor

[asfimport]

Ignacio Vera (@iverase) (migrated from JIRA)

In #10859 we have done the first step to move the API but still the methods #visitDocsIds and #visitDocValues are using the IntersectVisitor as an input.  Here I am proposing to introduce two functional interfaces DocIdsVisitor and DocValuesVisitor to use them as the input for those methods so the API would look like:

 

/**
 * Basic operations to read the KD-tree.
 *
 * `@lucene`.experimental
 */
public interface PointTree extends Cloneable {

  /** Clone, the current node becomes the root of the new tree. */
  PointTree clone();

  /**
   * Move to the first child node and return {`@code` true} upon success. Returns {`@code` false} for
   * leaf nodes and {`@code` true} otherwise.
   */
  boolean moveToChild() throws IOException;

  /**
   * Move to the next sibling node and return {`@code` true} upon success. Returns {`@code` false} if
   * the current node has no more siblings.
   */
  boolean moveToSibling() throws IOException;

  /**
   * Move to the parent node and return {`@code` true} upon success. Returns {`@code` false} for the
   * root node and {`@code` true} otherwise.
   */
  boolean moveToParent() throws IOException;

  /** Return the minimum packed value of the current node. */
  byte[] getMinPackedValue();

  /** Return the maximum packed value of the current node. */
  byte[] getMaxPackedValue();

  /** Return the number of points below the current node. */
  long size();

  /** Visit all the docs below the current node. */
  void visitDocIDs(DocIdsVisitor docIdsVisitor) throws IOException;

  /** Visit all the docs and values below the current node. */
  default void visitDocValues(DocValuesVisitor docValuesVisitor) throws IOException {
    visitDocValues((min, max) -> Relation.CELL_CROSSES_QUERY, docID -> {}, docValuesVisitor);
  }

  /**
   * Similar to {`@link` #visitDocValues(DocValuesVisitor)} but in this case it allows adding a
   * filter that works like {`@link` IntersectVisitor#compare(byte[], byte[])}.
   */
  void visitDocValues(
      BiFunction<byte[], byte[], Relation> compare,
      DocIdsVisitor docIdsVisitor,
      DocValuesVisitor docValuesVisitor)
      throws IOException;
}

/**
 * Collects all documents below a tree node by calling {`@link`
 * PointTree#visitDocIDs(DocIdsVisitor)}
 */
`@FunctionalInterface`
public interface DocIdsVisitor {
  /** Called for all documents below a tree node. */
  void visit(int docID) throws IOException;
}

/**
 * Collects all documents and values below a tree node by calling {`@link`
 * PointTree#visitDocValues(DocValuesVisitor)} (DocIdsVisitor)}
 */
`@FunctionalInterface`
public interface DocValuesVisitor {
  /** Called for all documents and values below a tree node. */
  void visit(int docID, byte[] packedValue) throws IOException;

  /**
   * Similar to {`@link` DocValuesVisitor#visit(int, byte[])} but in this case the packedValue can
   * have more than one docID associated to it. The provided iterator should not escape the scope
   * of this method so that implementations of PointValues are free to reuse it.
   */
  default void visit(DocIdSetIterator iterator, byte[] packedValue) throws IOException {
    int docID;
    while ((docID = iterator.nextDoc()) != DocIdSetIterator.NO_MORE_DOCS) {
      visit(docID, packedValue);
    }
  }
}

/**
 * We recurse the {`@link` PointTree}, using a provided instance of this to guide the recursion.
 *
 * `@lucene`.experimental
 */
public interface IntersectVisitor extends DocValuesVisitor, DocIdsVisitor {

  /**
   * Called for non-leaf cells to test how the cell relates to the query, to determine how to
   * further recurse down the tree.
   *
   * <ul>
   *   <li>{`@link` Relation#CELL_OUTSIDE_QUERY}: Stop recursing down the current branch of the
   *       tree.
   *   <li>{`@link` Relation#CELL_INSIDE_QUERY}: All nodes below the current node are visited using
   *       the underlying {`@link` DocIdsVisitor}. he consumer should generally blindly accept the
   *       docID.
   *   <li>{`@link` Relation#CELL_CROSSES_QUERY}: Keep recursing down the current branch of the
   *       tree. If the current node is a leaf, visit all docs and values usinng the underlying
   *       {`@link` DocValuesVisitor}. The consumer should scrutinize the packedValue to decide
   *       whether to accept it.
   * </ul>
   */
  Relation compare(byte[] minPackedValue, byte[] maxPackedValue);

  /** Notifies the caller that this many documents are about to be visited */
  default void grow(int count) {}
} 

 

Any thoughts?

[asfimport]

Adrien Grand (@jpountz) (migrated from JIRA)

We are losing the ability to take advantage of grow() with this proposal, aren't we? Also I'm a bit on the fence about the visitDocValues method that only takes a DocValuesVisitor since it encourages leveraging the tree in a slow way that reads all values and doc IDs?

[asfimport]

Ignacio Vera (@iverase) (migrated from JIRA)

We are losing the ability to take advantage of grow() with this proposal, aren't we?

What we are doing is to move the call to grow() to the intersects algorithm instead of expecting that PointTree implementations are calling it. I actually think this is better than the current implementation as I think trees should not have to do it (It is an intersects implementation detail). Anyway if we are not moving with this proposal we should document that PointTree implementations are expected to call grow() in #visitDocIDs and #visitDocValues methods?

 

I'm a bit on the fence about the visitDocValues method that only takes a DocValuesVisitor since it encourages leveraging the tree in a slow way that reads all values and doc IDs

I have the opposite feeling. If a user is using the programatic API, when it calls #visitDocValues, it wants to visit all doc and values. The method that takes the filter is only related to #intersects, and what we are trying to do is to encourage other uses? 

 

I think we should move the IntersectVisitor out as it seems to flag that the trees are only meant to do interacts like operations which is not true.

 

[asfimport]

Adrien Grand (@jpountz) (migrated from JIRA)

Thanks for explaining, I had misunderstood what the change that you were suggesting would do, I was confused about whether we were looking at new methods for PointTree or PointValues.

I think it makes sense to me, maybe make a PR for it (it looks like a small change, isn't it?) where it's going to be easier for me to see the implications?

[asfimport]

Ignacio Vera (@iverase) (migrated from JIRA)

PR is here: #486

[jpountz]

This has been implemented.