ZoieSegmentReader.java

package proj.zoie.api;

/**
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
import it.unimi.dsi.fastutil.ints.IntRBTreeSet;
import it.unimi.dsi.fastutil.longs.LongIterator;
import it.unimi.dsi.fastutil.longs.LongSet;

import java.io.IOException;
import java.util.Arrays;

import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.document.NumericDocValuesField;
import org.apache.lucene.index.AtomicReader;
import org.apache.lucene.index.FilterAtomicReader;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.NumericDocValues;
import org.apache.lucene.index.SegmentReader;
import org.apache.lucene.util.Bits;
import org.apache.lucene.util.BytesRef;

import proj.zoie.api.indexing.AbstractZoieIndexable;
import proj.zoie.api.indexing.IndexReaderDecorator;

public class ZoieSegmentReader<R extends IndexReader> extends FilterAtomicReader {
  public static final long DELETED_UID = Long.MIN_VALUE;
  private R _decoratedReader;
  private final IndexReaderDecorator<R> _decorator;
  private IntRBTreeSet _delDocIdSet = new IntRBTreeSet();
  private int[] _currentDelDocIds = null;
  private long[] _uidArray = null;
  private DocIDMapper _docIDMapper = null;

  public static void fillDocumentID(Document doc, long id) {
    Field uidField = new NumericDocValuesField(AbstractZoieIndexable.DOCUMENT_ID_PAYLOAD_FIELD, id);
    doc.add(uidField);
  }

  public ZoieSegmentReader(AtomicReader in, IndexReaderDecorator<R> decorator) throws IOException {
    super(in);
    if (!(in instanceof SegmentReader)) {
      throw new IllegalStateException("ZoieSegmentReader can only be constucted from "
          + SegmentReader.class);
    }
    init(in);
    _decorator = decorator;
    _decoratedReader = (_decorator == null ? null : _decorator.decorate(this));
  }

  private void init(AtomicReader reader) throws IOException {
    int maxDoc = reader.maxDoc();
    _uidArray = new long[maxDoc];
    NumericDocValues uidValues = reader
        .getNumericDocValues(AbstractZoieIndexable.DOCUMENT_ID_PAYLOAD_FIELD);
    Bits liveDocs = reader.getLiveDocs();
    for (int i = 0; i < maxDoc; ++i) {
      if (liveDocs != null && !liveDocs.get(i)) {
        _uidArray[i] = ZoieSegmentReader.DELETED_UID;
        continue;
      }
      _uidArray[i] = uidValues.get(i);
    }
  }

  /**
   * make exact shallow copy for duplication. The decorated reader is also shallow copied.
   * @param copyFrom
   * @param innerReader
   * @throws IOException
   */
  ZoieSegmentReader(ZoieSegmentReader<R> copyFrom, AtomicReader innerReader) throws IOException {
    super(innerReader);
    _docIDMapper = copyFrom._docIDMapper;
    _decorator = copyFrom._decorator;
    _delDocIdSet = copyFrom._delDocIdSet;
    _currentDelDocIds = copyFrom._currentDelDocIds;
    _uidArray = copyFrom._uidArray;

    if (copyFrom._decorator == null) {
      _decoratedReader = null;
    } else {
      _decoratedReader = copyFrom._decorator.redecorate(copyFrom._decoratedReader, this);
    }
  }

  /**
   * makes exact shallow copy of a given ZoieSegmentReader
   * @throws IOException
   */
  public ZoieSegmentReader<R> copy() throws IOException {
    return new ZoieSegmentReader<R>(this, this.in);
  }

  public AtomicReader getInnerReader() {
    return in;
  }

  @Override
  public Bits getLiveDocs() {
    ensureOpen();
    return new Bits() {
      @Override
      public boolean get(int index) {
        int[] delSet = _currentDelDocIds;
        if (delSet != null && Arrays.binarySearch(delSet, index) >= 0) {
          return false;
        }
        Bits liveDocs = in.getLiveDocs();
        if (liveDocs == null) {
          return true;
        }
        return liveDocs.get(index);
      }

      @Override
      public int length() {
        return in.getLiveDocs().length();
      }
    };
  }

  public void markDeletes(LongSet delDocs, LongSet deletedUIDs) {
    LongIterator iter = delDocs.iterator();
    IntRBTreeSet delDocIdSet = _delDocIdSet;

    while (iter.hasNext()) {
      long uid = iter.nextLong();
      if (ZoieSegmentReader.DELETED_UID != uid) {
        int docid = _docIDMapper.getDocID(uid);
        if (docid != DocIDMapper.NOT_FOUND) {
          delDocIdSet.add(docid);
          deletedUIDs.add(uid);
        }
      }
    }
  }

  public void commitDeletes() {
    _currentDelDocIds = _delDocIdSet.toIntArray();
  }

  public R getDecoratedReader() {
    return _decoratedReader;
  }

  public BytesRef getStoredValue(int docid) throws IOException {
    Document doc = in.document(docid);
    if (doc != null) {
      return doc.getBinaryValue(AbstractZoieIndexable.DOCUMENT_STORE_FIELD);
    }
    return null;
  }

  public DocIDMapper getDocIDMapper() {
    return _docIDMapper;
  }

  public long[] getUIDArray() {
    return _uidArray;
  }

  public void setDocIDMapper(DocIDMapper docIDMapper) {
    _docIDMapper = docIDMapper;
  }

  public long getUID(int docid) {
    return _uidArray[docid];
  }

  public boolean isDeleted(int docid) {
    int[] delSet = _currentDelDocIds;
    if (delSet != null && Arrays.binarySearch(delSet, docid) >= 0) {
      return true;
    }
    Bits liveDocs = in.getLiveDocs();
    if (liveDocs == null) {
      return false;
    }
    return !liveDocs.get(docid);
  }

  public boolean isDeletedInMask(int docid) {
    int[] delSet = _currentDelDocIds;
    if (delSet != null && Arrays.binarySearch(delSet, docid) >= 0) {
      return true;
    }
    return false;
  }

  public String getSegmentName() {
    return ((SegmentReader) in).getSegmentName();
  }

  @Override
  public int numDocs() {
    if (_currentDelDocIds != null) {
      return super.numDocs() - _currentDelDocIds.length;
    } else {
      return super.numDocs();
    }
  }

}