io.mindmaps.graql.internal.analytics.DegreeAndPersistVertexProgram Maven / Gradle / Ivy

Go to download
/*
 * MindmapsDB - A Distributed Semantic Database
 * Copyright (C) 2016  Mindmaps Research Ltd
 *
 * MindmapsDB 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.
 *
 * MindmapsDB 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 MindmapsDB. If not, see .
 */

package io.mindmaps.graql.internal.analytics;

import io.mindmaps.util.Schema;
import org.apache.tinkerpop.gremlin.process.computer.GraphComputer;
import org.apache.tinkerpop.gremlin.process.computer.Memory;
import org.apache.tinkerpop.gremlin.process.computer.Messenger;
import org.apache.tinkerpop.gremlin.structure.Vertex;
import org.apache.tinkerpop.gremlin.util.iterator.IteratorUtils;

import java.util.Collections;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Set;

class DegreeAndPersistVertexProgram extends MindmapsVertexProgram {

    private static final String KEYSPACE_KEY = "keyspace";

    private BulkResourceMutate bulkResourceMutate;

    public DegreeAndPersistVertexProgram() {
    }

    public DegreeAndPersistVertexProgram(String keySpace, Set types) {
        persistentProperties.put(KEYSPACE_KEY, keySpace);
        selectedTypes = types;
    }

    @Override
    public GraphComputer.Persist getPreferredPersist() {
        return GraphComputer.Persist.NOTHING;
    }

    @Override
    public void safeExecute(final Vertex vertex, Messenger messenger, final Memory memory) {
        switch (memory.getIteration()) {
            case 0:
                if (selectedTypes.contains(Utility.getVertexType(vertex)) && !Utility.isAnalyticsElement(vertex)) {
                    String type = vertex.value(Schema.ConceptProperty.BASE_TYPE.name());
                    if (type.equals(Schema.BaseType.ENTITY.name()) || type.equals(Schema.BaseType.RESOURCE.name())) {
                        messenger.sendMessage(countMessageScopeIn, 1L);
                    } else if (type.equals(Schema.BaseType.RELATION.name())) {
                        messenger.sendMessage(countMessageScopeOut, -1L);
                        messenger.sendMessage(countMessageScopeIn, 1L);
                    }
                }
                break;
            case 1:
                if (vertex.value(Schema.ConceptProperty.BASE_TYPE.name()).equals(Schema.BaseType.CASTING.name())) {
                    boolean hasRolePlayer = false;
                    long assertionCount = 0;
                    Iterator iterator = messenger.receiveMessages();
                    while (iterator.hasNext()) {
                        long message = iterator.next();
                        if (message < 0) assertionCount++;
                        else hasRolePlayer = true;
                    }
                    if (hasRolePlayer) {
                        messenger.sendMessage(countMessageScopeIn, 1L);
                        messenger.sendMessage(countMessageScopeOut, assertionCount);
                    }
                }
                break;
            case 2:
                if (!Utility.isAnalyticsElement(vertex) && selectedTypes.contains(Utility.getVertexType(vertex))) {
                    long edgeCount = IteratorUtils.reduce(messenger.receiveMessages(), 0L, (a, b) -> a + b);
                    bulkResourceMutate.putValue(vertex, edgeCount);
                }
                break;
        }
    }

    @Override
    public void workerIterationStart(Memory memory) {
        bulkResourceMutate = new BulkResourceMutate((String) persistentProperties.get(KEYSPACE_KEY));
    }

    @Override
    public void workerIterationEnd(Memory memory) {
        bulkResourceMutate.flush();
    }

    @Override
    public boolean terminate(final Memory memory) {
        return memory.getIteration() == 2;
    }

}