lucene/core/src/test/org/apache/lucene/index/TestMultiLevelSkipList.java

/*
 * 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.
 */
package org.apache.lucene.index;


import java.io.IOException;
import java.util.concurrent.atomic.AtomicInteger;

import org.apache.lucene.analysis.*;
import org.apache.lucene.analysis.tokenattributes.PayloadAttribute;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.store.Directory;
import org.apache.lucene.store.IOContext;
import org.apache.lucene.store.IndexInput;
import org.apache.lucene.store.MockDirectoryWrapper;
import org.apache.lucene.store.RAMDirectory;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.LuceneTestCase;
import org.apache.lucene.util.TestUtil;
import org.junit.Before;

/**
 * This testcase tests whether multi-level skipping is being used
 * to reduce I/O while skipping through posting lists.
 * 
 * Skipping in general is already covered by several other
 * testcases.
 * 
 */
public class TestMultiLevelSkipList extends LuceneTestCase {
  
  class CountingRAMDirectory extends MockDirectoryWrapper {
    public CountingRAMDirectory(Directory delegate) {
      super(random(), delegate);
    }

    @Override
    public IndexInput openInput(String fileName, IOContext context) throws IOException {
      IndexInput in = super.openInput(fileName, context);
      if (fileName.endsWith(".frq"))
        in = new CountingStream(in);
      return in;
    }
  }
  
  @Override
  @Before
  public void setUp() throws Exception {
    super.setUp();
    counter = 0;
  }

  public void testSimpleSkip() throws IOException {
    Directory dir = new CountingRAMDirectory(new RAMDirectory());
    IndexWriter writer = new IndexWriter(dir, newIndexWriterConfig(new PayloadAnalyzer())
                                                .setCodec(TestUtil.alwaysPostingsFormat(TestUtil.getDefaultPostingsFormat()))
                                                .setMergePolicy(newLogMergePolicy()));
    Term term = new Term("test", "a");
    for (int i = 0; i < 5000; i++) {
      Document d1 = new Document();
      d1.add(newTextField(term.field(), term.text(), Field.Store.NO));
      writer.addDocument(d1);
    }
    writer.commit();
    writer.forceMerge(1);
    writer.close();

    LeafReader reader = getOnlyLeafReader(DirectoryReader.open(dir));
    
    for (int i = 0; i < 2; i++) {
      counter = 0;
      PostingsEnum tp = reader.postings(term, PostingsEnum.ALL);
      checkSkipTo(tp, 14, 185); // no skips
      checkSkipTo(tp, 17, 190); // one skip on level 0
      checkSkipTo(tp, 287, 200); // one skip on level 1, two on level 0
    
      // this test would fail if we had only one skip level,
      // because than more bytes would be read from the freqStream
      checkSkipTo(tp, 4800, 250);// one skip on level 2
    }
  }

  public void checkSkipTo(PostingsEnum tp, int target, int maxCounter) throws IOException {
    tp.advance(target);
    if (maxCounter < counter) {
      fail("Too many bytes read: " + counter + " vs " + maxCounter);
    }

    assertEquals("Wrong document " + tp.docID() + " after skipTo target " + target, target, tp.docID());
    assertEquals("Frequency is not 1: " + tp.freq(), 1,tp.freq());
    tp.nextPosition();
    BytesRef b = tp.getPayload();
    assertEquals(1, b.length);
    assertEquals("Wrong payload for the target " + target + ": " + b.bytes[b.offset], (byte) target, b.bytes[b.offset]);
  }

  private static class PayloadAnalyzer extends Analyzer {
    private final AtomicInteger payloadCount = new AtomicInteger(-1);
    @Override
    public TokenStreamComponents createComponents(String fieldName) {
      Tokenizer tokenizer = new MockTokenizer(MockTokenizer.WHITESPACE, true);
      return new TokenStreamComponents(tokenizer, new PayloadFilter(payloadCount, tokenizer));
    }

  }

  private static class PayloadFilter extends TokenFilter {
    
    PayloadAttribute payloadAtt;
    private AtomicInteger payloadCount;
    
    protected PayloadFilter(AtomicInteger payloadCount , TokenStream input) {
      super(input);
      this.payloadCount = payloadCount;
      payloadAtt = addAttribute(PayloadAttribute.class);
    }

    @Override
    public boolean incrementToken() throws IOException {
      boolean hasNext = input.incrementToken();
      if (hasNext) {
        payloadAtt.setPayload(new BytesRef(new byte[] { (byte) payloadCount.incrementAndGet() }));
      } 
      return hasNext;
    }

  }

  private int counter = 0;

  // Simply extends IndexInput in a way that we are able to count the number
  // of bytes read
  class CountingStream extends IndexInput {
    private IndexInput input;

    CountingStream(IndexInput input) {
      super("CountingStream(" + input + ")");
      this.input = input;
    }

    @Override
    public byte readByte() throws IOException {
      TestMultiLevelSkipList.this.counter++;
      return this.input.readByte();
    }

    @Override
    public void readBytes(byte[] b, int offset, int len) throws IOException {
      TestMultiLevelSkipList.this.counter += len;
      this.input.readBytes(b, offset, len);
    }

    @Override
    public void close() throws IOException {
      this.input.close();
    }

    @Override
    public long getFilePointer() {
      return this.input.getFilePointer();
    }

    @Override
    public void seek(long pos) throws IOException {
      this.input.seek(pos);
    }

    @Override
    public long length() {
      return this.input.length();
    }

    @Override
    public CountingStream clone() {
      return new CountingStream(this.input.clone());
    }

    @Override
    public IndexInput slice(String sliceDescription, long offset, long length) throws IOException {
      return new CountingStream(this.input.slice(sliceDescription, offset, length));
    }
  }
}