org.neo4j.index.impl.lucene.legacy.LuceneBatchInserterIndex.java Source code

Java tutorial

Introduction

Here is the source code for org.neo4j.index.impl.lucene.legacy.LuceneBatchInserterIndex.java

Source

/*
 * Copyright (c) 2002-2016 "Neo Technology,"
 * Network Engine for Objects in Lund AB [http://neotechnology.com]
 *
 * This file is part of Neo4j.
 *
 * Neo4j is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package org.neo4j.index.impl.lucene.legacy;

import org.apache.lucene.document.Document;
import org.apache.lucene.index.CorruptIndexException;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.index.IndexWriterConfig;
import org.apache.lucene.index.IndexableField;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.SearcherFactory;
import org.apache.lucene.search.SearcherManager;
import org.apache.lucene.search.Sort;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.store.Directory;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;

import org.neo4j.collection.primitive.PrimitiveLongCollections;
import org.neo4j.graphdb.ResourceIterator;
import org.neo4j.graphdb.index.IndexHits;
import org.neo4j.helpers.collection.LruCache;
import org.neo4j.index.lucene.ValueContext;
import org.neo4j.io.IOUtils;
import org.neo4j.kernel.api.LegacyIndexHits;
import org.neo4j.kernel.api.impl.index.collector.DocValuesCollector;
import org.neo4j.kernel.impl.index.IndexEntityType;
import org.neo4j.kernel.impl.util.IoPrimitiveUtils;
import org.neo4j.unsafe.batchinsert.BatchInserterIndex;

class LuceneBatchInserterIndex implements BatchInserterIndex {
    private final IndexIdentifier identifier;
    private final IndexType type;

    private IndexWriter writer;
    private SearcherManager searcherManager;
    private final boolean createdNow;
    private Map<String, LruCache<String, Collection<EntityId>>> cache;
    private int updateCount;
    private final int commitBatchSize = 500000;
    private final RelationshipLookup relationshipLookup;

    interface RelationshipLookup {
        EntityId lookup(long id);
    }

    LuceneBatchInserterIndex(File dbStoreDir, IndexIdentifier identifier, Map<String, String> config,
            RelationshipLookup relationshipLookup) {
        File storeDir = getStoreDir(dbStoreDir);
        this.createdNow = !LuceneDataSource.getFileDirectory(storeDir, identifier).exists();
        this.identifier = identifier;
        this.type = IndexType.getIndexType(config);
        this.relationshipLookup = relationshipLookup;
        this.writer = instantiateWriter(storeDir);
        this.searcherManager = instantiateSearcherManager(writer);
    }

    @Override
    public void add(long id, Map<String, Object> properties) {
        try {
            Document document = IndexType.newDocument(entityId(id));
            for (Map.Entry<String, Object> entry : properties.entrySet()) {
                String key = entry.getKey();
                Object value = entry.getValue();
                addSingleProperty(id, document, key, value);
            }
            writer.addDocument(document);
            if (++updateCount == commitBatchSize) {
                writer.commit();
                updateCount = 0;
            }
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    private EntityId entityId(long id) {
        if (identifier.entityType == IndexEntityType.Node) {
            return new EntityId.IdData(id);
        }

        return relationshipLookup.lookup(id);
    }

    private void addSingleProperty(long entityId, Document document, String key, Object value) {
        for (Object oneValue : IoPrimitiveUtils.asArray(value)) {
            boolean isValueContext = oneValue instanceof ValueContext;
            oneValue = isValueContext ? ((ValueContext) oneValue).getCorrectValue() : oneValue.toString();
            type.addToDocument(document, key, oneValue);
            if (createdNow) {
                // If we know that the index was created this session
                // then we can go ahead and add stuff to the cache directly
                // when adding to the index.
                addToCache(entityId, key, oneValue);
            }
        }
    }

    private void addToCache(long entityId, String key, Object value) {
        if (this.cache == null) {
            return;
        }

        String valueAsString = value.toString();
        LruCache<String, Collection<EntityId>> cache = this.cache.get(key);
        if (cache != null) {
            Collection<EntityId> ids = cache.get(valueAsString);
            if (ids == null) {
                ids = new HashSet<>();
                cache.put(valueAsString, ids);
            }
            ids.add(new EntityId.IdData(entityId));
        }
    }

    private void addToCache(Collection<EntityId> ids, String key, Object value) {
        if (this.cache == null) {
            return;
        }

        String valueAsString = value.toString();
        LruCache<String, Collection<EntityId>> cache = this.cache.get(key);
        if (cache != null) {
            cache.put(valueAsString, ids);
        }
    }

    private LegacyIndexHits getFromCache(String key, Object value) {
        if (this.cache == null) {
            return null;
        }

        String valueAsString = value.toString();
        LruCache<String, Collection<EntityId>> cache = this.cache.get(key);
        if (cache != null) {
            Collection<EntityId> ids = cache.get(valueAsString);
            if (ids != null) {
                return new ConstantScoreIterator(ids, Float.NaN);
            }
        }
        return null;
    }

    @Override
    public void updateOrAdd(long entityId, Map<String, Object> properties) {
        try {
            removeFromCache(entityId);
            writer.deleteDocuments(type.idTermQuery(entityId));
            add(entityId, properties);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    private void removeFromCache(long entityId) throws IOException, CorruptIndexException {
        IndexSearcher searcher = searcherManager.acquire();
        try {
            Query query = type.idTermQuery(entityId);
            TopDocs docs = searcher.search(query, 1);
            if (docs.totalHits > 0) {
                Document document = searcher.doc(docs.scoreDocs[0].doc);
                for (IndexableField field : document.getFields()) {
                    String key = field.name();
                    Object value = field.stringValue();
                    removeFromCache(entityId, key, value);
                }
            }
        } finally {
            searcherManager.release(searcher);
        }
    }

    private void removeFromCache(long entityId, String key, Object value) {
        if (this.cache == null) {
            return;
        }

        String valueAsString = value.toString();
        LruCache<String, Collection<EntityId>> cache = this.cache.get(key);
        if (cache != null) {
            Collection<EntityId> ids = cache.get(valueAsString);
            if (ids != null) {
                ids.remove(new EntityId.IdData(entityId));
            }
        }
    }

    private IndexWriter instantiateWriter(File folder) {
        Directory dir = null;
        try {
            dir = LuceneDataSource.getDirectory(folder, identifier);
            IndexWriterConfig writerConfig = new IndexWriterConfig(type.analyzer);
            writerConfig.setRAMBufferSizeMB(determineGoodBufferSize(writerConfig.getRAMBufferSizeMB()));
            return new IndexWriter(dir, writerConfig);
        } catch (IOException e) {
            IOUtils.closeAllSilently(dir);
            throw new RuntimeException(e);
        }
    }

    private double determineGoodBufferSize(double atLeast) {
        double heapHint = Runtime.getRuntime().maxMemory() / (1024 * 1024 * 14);
        double result = Math.max(atLeast, heapHint);
        return Math.min(result, 700);
    }

    private static SearcherManager instantiateSearcherManager(IndexWriter writer) {
        try {
            return new SearcherManager(writer, true, new SearcherFactory());
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    private void closeSearcher() {
        try {
            if (searcherManager != null) {
                this.searcherManager.close();
            }
        } catch (IOException e) {
            throw new RuntimeException(e);
        } finally {
            this.searcherManager = null;
        }
    }

    private void closeWriter() {
        try {
            if (this.writer != null) {
                this.writer.close();
            }
        } catch (IOException e) {
            throw new RuntimeException(e);
        } finally {
            this.writer = null;
        }
    }

    private IndexHits<Long> query(Query query, final String key, final Object value) {
        IndexSearcher searcher;
        try {
            searcher = searcherManager.acquire();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        try {
            DocValuesCollector collector = new DocValuesCollector(true);
            searcher.search(query, collector);
            IndexHits<Document> result = collector.getIndexHits(Sort.RELEVANCE);
            LegacyIndexHits primitiveHits = null;
            if (key == null || this.cache == null || !this.cache.containsKey(key)) {
                primitiveHits = new DocToIdIterator(result, Collections.<EntityId>emptyList(), null,
                        PrimitiveLongCollections.emptySet());
            } else {
                primitiveHits = new DocToIdIterator(result, Collections.<EntityId>emptyList(), null,
                        PrimitiveLongCollections.emptySet()) {
                    private final Collection<EntityId> ids = new ArrayList<>();

                    @Override
                    protected boolean fetchNext() {
                        if (super.fetchNext()) {
                            ids.add(new EntityId.IdData(next));
                            return true;
                        }
                        addToCache(ids, key, value);
                        return false;
                    }
                };
            }
            return wrapIndexHits(primitiveHits);
        } catch (IOException e) {
            throw new RuntimeException(e);
        } finally {
            try {
                searcherManager.release(searcher);
            } catch (IOException ignore) {
            }
        }
    }

    private IndexHits<Long> wrapIndexHits(final LegacyIndexHits ids) {
        return new IndexHits<Long>() {
            @Override
            public boolean hasNext() {
                return ids.hasNext();
            }

            @Override
            public Long next() {
                return ids.next();
            }

            @Override
            public void remove() {
                throw new UnsupportedOperationException();
            }

            @Override
            public ResourceIterator<Long> iterator() {
                return this;
            }

            @Override
            public int size() {
                return ids.size();
            }

            @Override
            public void close() {
                ids.close();
            }

            @Override
            public Long getSingle() {
                try {
                    long singleId = PrimitiveLongCollections.single(ids, -1L);
                    return singleId == -1 ? null : singleId;
                } finally {
                    close();
                }
            }

            @Override
            public float currentScore() {
                return 0;
            }
        };
    }

    @Override
    public IndexHits<Long> get(String key, Object value) {
        LegacyIndexHits cached = getFromCache(key, value);
        return cached != null ? wrapIndexHits(cached) : query(type.get(key, value), key, value);
    }

    @Override
    public IndexHits<Long> query(String key, Object queryOrQueryObject) {
        return query(type.query(key, queryOrQueryObject, null), null, null);
    }

    @Override
    public IndexHits<Long> query(Object queryOrQueryObject) {
        return query(type.query(null, queryOrQueryObject, null), null, null);
    }

    public void shutdown() {
        closeSearcher();
        closeWriter();
    }

    private File getStoreDir(File dbStoreDir) {
        File dir = new File(dbStoreDir, "index");
        if (!dir.exists() && !dir.mkdirs()) {
            throw new RuntimeException(
                    "Unable to create directory path[" + dir.getAbsolutePath() + "] for Neo4j store.");
        }
        return dir;
    }

    @Override
    public void flush() {
        try {
            searcherManager.maybeRefreshBlocking();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public void setCacheCapacity(String key, int size) {
        if (this.cache == null) {
            this.cache = new HashMap<>();
        }
        LruCache<String, Collection<EntityId>> cache = this.cache.get(key);
        if (cache != null) {
            cache.resize(size);
        } else {
            cache = new LruCache<>("Batch inserter cache for " + key, size);
            this.cache.put(key, cache);
        }
    }
}