cn.hbu.cs.esearch.core.EsearchMultiReader.java Source code

Java tutorial

Introduction

Here is the source code for cn.hbu.cs.esearch.core.EsearchMultiReader.java

Source

/*
 * 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 cn.hbu.cs.esearch.core;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicLong;

import org.apache.lucene.index.AtomicReader;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.index.FilterDirectoryReader;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.SegmentReader;
import org.apache.lucene.util.BytesRef;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import cn.hbu.cs.esearch.document.DocIDMapper;
import cn.hbu.cs.esearch.index.IndexReaderDecorator;
import it.unimi.dsi.fastutil.longs.LongSet;

public class EsearchMultiReader<R extends IndexReader> extends FilterDirectoryReader {
    public static final Logger LOGGER = LoggerFactory.getLogger(EsearchMultiReader.class);
    private final Map<String, EsearchSegmentReader<R>> _readerMap;
    private final List<EsearchSegmentReader<R>> _subEsearchReaders;
    private List<R> _decoratedReaders;
    private final IndexReaderDecorator<R> _decorator;
    private DocIDMapper _docIDMapper;

    public EsearchMultiReader(DirectoryReader in, IndexReaderDecorator<R> decorator) throws IOException {
        this(in, decorator, new EsearchSubReaderWrapper<R>(decorator));
    }

    @SuppressWarnings("unchecked")
    private EsearchMultiReader(DirectoryReader in, IndexReaderDecorator<R> decorator,
            EsearchSubReaderWrapper<R> wrapper) throws IOException {
        super(in, wrapper);
        _subEsearchReaders = (List<EsearchSegmentReader<R>>) getSequentialSubReaders();
        _decorator = decorator;
        _readerMap = new HashMap<String, EsearchSegmentReader<R>>();
        _decoratedReaders = null;
        init();
    }

    private final AtomicLong zoieRefCounter = new AtomicLong(1);

    public void incEsearchRef() {
        zoieRefCounter.incrementAndGet();
    }

    public void decEsearchRef() {
        long refCount = zoieRefCounter.decrementAndGet();
        if (refCount < 0) {
            LOGGER.warn("refCount should never be lower than 0");
        }
        if (refCount == 0) {
            try {
                in.decRef();
            } catch (IOException e) {
                LOGGER.error("decEsearchRef exception, ", e);
            }
        }
    }

    public int getInnerRefCount() {
        return in.getRefCount();
    }

    public int getSubReaderBase(int idx) {
        return readerBase(idx);
    }

    public DocIDMapper getDocIDMapper() {
        return _docIDMapper;
    }

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

    public BytesRef getStoredValue(long uid) throws IOException {
        int docid = _docIDMapper.getDocID(uid);
        if (docid < 0) {
            return null;
        }
        int idx = readerIndex(docid);
        if (idx < 0) {
            return null;
        }
        EsearchSegmentReader<R> subReader = _subEsearchReaders.get(idx);
        return subReader.getStoredValue(docid - readerBase(idx));
    }

    private void init() throws IOException {
        for (EsearchSegmentReader<R> subReader : _subEsearchReaders) {
            String segmentName = subReader.getSegmentName();
            _readerMap.put(segmentName, subReader);
        }

        ArrayList<R> decoratedList = new ArrayList<R>(_subEsearchReaders.size());
        for (EsearchSegmentReader<R> subReader : _subEsearchReaders) {
            R decoratedReader = subReader.getDecoratedReader();
            decoratedList.add(decoratedReader);
        }
        _decoratedReaders = decoratedList;
    }

    @SuppressWarnings("unchecked")
    public EsearchSegmentReader<R>[] getSubReaders() {
        return (_subEsearchReaders.toArray(new EsearchSegmentReader[_subEsearchReaders.size()]));
    }

    public void markDeletes(LongSet delDocs, LongSet deletedUIDs) {
        EsearchSegmentReader<R>[] subReaders = getSubReaders();
        if (subReaders != null && subReaders.length > 0) {
            for (EsearchSegmentReader<R> subReader : subReaders) {
                subReader.markDeletes(delDocs, deletedUIDs);
            }
        }
    }

    public void commitDeletes() {
        EsearchSegmentReader<R>[] subReaders = getSubReaders();
        if (subReaders != null && subReaders.length > 0) {
            for (EsearchSegmentReader<R> subReader : subReaders) {
                subReader.commitDeletes();
            }
        }
    }

    public List<R> getDecoratedReaders() throws IOException {
        return _decoratedReaders;
    }

    public static <R extends IndexReader> List<R> extractDecoratedReaders(List<EsearchMultiReader<R>> readerList)
            throws IOException {
        LinkedList<R> retList = new LinkedList<R>();
        for (EsearchMultiReader<R> reader : readerList) {
            retList.addAll(reader.getDecoratedReaders());
        }
        return retList;
    }

    public boolean isDeleted(int docid) {
        int idx = readerIndex(docid);
        EsearchSegmentReader<R> subReader = _subEsearchReaders.get(idx);
        return subReader.isDeleted(docid - readerBase(idx));
    }

    public EsearchMultiReader<R> reopen() throws IOException {
        long t0 = System.currentTimeMillis();
        DirectoryReader inner = DirectoryReader.openIfChanged(in);
        if (inner == null) {
            t0 = System.currentTimeMillis() - t0;
            if (t0 > 1000) {
                LOGGER.info("reopen returns in " + t0 + "ms without change");
            } else {
                if (LOGGER.isDebugEnabled()) {
                    LOGGER.debug("reopen returns in " + t0 + "ms without change");
                }
            }
            return this;
        }

        EsearchMultiReader<R> ret = new EsearchMultiReader<R>(inner, _decorator,
                new EsearchSubReaderWrapper<R>(_decorator, _readerMap));
        t0 = System.currentTimeMillis() - t0;
        if (t0 > 1000) {
            LOGGER.info("reopen returns in " + t0 + "ms with change");
        } else {
            if (LOGGER.isDebugEnabled()) {
                LOGGER.debug("reopen returns in " + t0 + "ms with change");
            }
        }
        return ret;
    }

    /**
     * makes exact shallow copy of a given EsearchMultiReader
     *
     * @throws java.io.IOException
     */
    public EsearchMultiReader<R> copy() throws IOException {
        // increase DirectoryReader reference counter
        this.in.incRef();
        EsearchMultiReader<R> ret = new EsearchMultiReader<R>(this.in, this._decorator,
                new EsearchSubReaderWrapper<R>(this._decorator, this._readerMap));
        ret._docIDMapper = this._docIDMapper;
        return ret;
    }

    @Override
    protected DirectoryReader doWrapDirectoryReader(DirectoryReader in) {
        return in;
    }

    public static class EsearchSubReaderWrapper<R extends IndexReader> extends SubReaderWrapper {
        private final IndexReaderDecorator<R> _decorator;
        private final Map<String, EsearchSegmentReader<R>> _readerMap;

        /**
         * Constructor
         */
        public EsearchSubReaderWrapper(IndexReaderDecorator<R> decorator) {
            this(decorator, null);
        }

        public EsearchSubReaderWrapper(IndexReaderDecorator<R> decorator,
                Map<String, EsearchSegmentReader<R>> readerMap) {
            _decorator = decorator;
            _readerMap = readerMap;
        }

        @Override
        public AtomicReader wrap(AtomicReader reader) {
            if (!(reader instanceof SegmentReader)) {
                throw new IllegalStateException("reader not insance of " + SegmentReader.class);
            }

            try {
                if (_readerMap != null && !_readerMap.isEmpty()) {
                    SegmentReader sr = (SegmentReader) reader;
                    String segmentName = sr.getSegmentName();
                    EsearchSegmentReader<R> zoieSegmentReader = _readerMap.get(segmentName);
                    if (zoieSegmentReader != null && zoieSegmentReader.getInnerReader() == sr) {
                        return new EsearchSegmentReader<R>(zoieSegmentReader, sr);
                    }
                }
                return new EsearchSegmentReader<R>(reader, _decorator);
            } catch (IOException e) {
                e.printStackTrace();
                throw new RuntimeException(e);
            }
        }
    }
}