Sqlg query Paths with condition on a previous Edge

SqlgTest.java

package unit;

import org.apache.commons.configuration.BaseConfiguration;
import org.apache.commons.configuration.Configuration;
import org.apache.tinkerpop.gremlin.process.traversal.Operator;
import org.apache.tinkerpop.gremlin.process.traversal.P;
import org.apache.tinkerpop.gremlin.process.traversal.Traversal;
import org.apache.tinkerpop.gremlin.process.traversal.Traverser;
import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.GraphTraversal;
import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.GraphTraversalSource;
import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.__;
import org.apache.tinkerpop.gremlin.structure.*;
import org.apache.tinkerpop.gremlin.structure.util.GraphFactory;
import org.junit.Assert;
import org.junit.Ignore;
import org.junit.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.umlg.sqlg.structure.SqlgGraph;

import java.util.Map;
import java.util.function.Function;
import java.util.function.Predicate;

import static org.apache.tinkerpop.gremlin.process.traversal.Pop.last;
import static org.apache.tinkerpop.gremlin.process.traversal.Scope.local;
//Tinkerpop 3.2.3 release
//Sqlg 1.3.2 release
// * Find All Paths with max length 20
// * only follow edges with the same speed values
// * if there are multiple edges to the same vertex, pick the one with the smallest difference between departure time (from outV) and arrival time (from inV)
// * ignore negative differences
// * Ej->Vj->Ej+1  - condition Ej.arrTime <= Ej+1.depTime  and  (Ej+1.depTime - Ej.arrTime) is smallest value in all edges that go to Vj+1 
// * there two version of query
public class SqlgTest {
	private static final Logger LOG = LoggerFactory.getLogger(SqlgTest.class);

	@Test
	public void sqlgTest() {
		Configuration config = new BaseConfiguration();
		config.setProperty("jdbc.url", "jdbc:postgresql://localhost:5432/sqlg_test");
		config.setProperty("jdbc.username", "user");
		config.setProperty("jdbc.password", "password");
		SqlgGraph graph = SqlgGraph.open(config);
		GraphTraversalSource g = graph.traversal();
		graph.tx().open();
		loadData(graph);
		graph.tx().commit();
		GraphTraversal gp = query1(g);
		checkResult(gp);
		gp = query2(g);
		checkResult(gp);
	}

	@Test
	public void tinkerGraphTest() {
		Configuration config = new BaseConfiguration();
		config.setProperty("gremlin.graph", "org.apache.tinkerpop.gremlin.tinkergraph.structure.TinkerGraph");
		Graph graph = GraphFactory.open(config);
		GraphTraversalSource g = graph.traversal();
		loadData(graph);
		GraphTraversal gp = query1(g);
		checkResult(gp);
		gp = query2(g);
		checkResult(gp);
	}


	private void checkResult(GraphTraversal gp) {
		int count = 0;
		while (gp.hasNext()) {
			LOG.info(gp.next().toString());
			count++;
		}
		Assert.assertEquals(8, count);
	}

	private void loadData(Graph graph) {
		Vertex v0 = graph.addVertex("code", "0");
		Vertex v1 = graph.addVertex("code", "1");
		Vertex v2 = graph.addVertex("code", "2");
		Vertex v3 = graph.addVertex("code", "3");

		v0.addEdge("tsw", v1, "speed", "1", "arrTime", 10l, "depTime", 5l);
		v1.addEdge("tsw", v2, "speed", "1", "arrTime", 15l, "depTime", 9l); //must be ignored in longest path
		v1.addEdge("tsw", v2, "speed", "1", "arrTime", 20l, "depTime", 17l); //must be used in longest path
		v2.addEdge("tsw", v3, "speed", "1", "arrTime", 30l, "depTime", 25l);
	}

	private GraphTraversal query1(GraphTraversalSource g) {
		Function timeAtWarehouse = new Function() {
			@Override
			public Long apply(Object o) {
				Map m = (Map) o;
				Long dt = ((Edge) (m.get("curr"))).value("depTime");
				Long at = ((Edge) (m.get("prev"))).value("arrTime");
				return (dt - at) >= 0 ? (dt - at) : Long.MAX_VALUE;
			}
		};

		Predicate checkNegativeTime = new Predicate() {
			@Override
			public boolean test(Object o) {
				Map m = (Map) (((Traverser) o).get());
				Long dt = ((Edge) (m.get("curr"))).value("depTime");
				Long at = ((Edge) (m.get("prev"))).value("arrTime");
				return (dt - at) >= 0;
			}
		};

		return g.V().outE("tsw").as("e").inV().emit().repeat(
				__.flatMap(
						__.outE("tsw").filter(__.as("edge").select(last, "e").where(P.eq("edge")).by("speed")).
								group().by(__.inV()).by(__.project("curr", "prev").by().by(__.select(last, "e")).fold()).select(Column.values).unfold().
								order(local).by(timeAtWarehouse).
								limit(local, 1).
								filter(checkNegativeTime).
								select("curr")
				).as("e").inV().simplePath()
		).times(20).map(__.union((Traversal) __.select(last, "e").by("speed"), (Traversal) __.path().by(T.id).by(__.valueMap("arrTime", "depTime"))).fold());
	}

	private GraphTraversal query2(GraphTraversalSource g) {
		return g.withSack(0).V().outE("tsw").as("e").inV().emit().repeat(
				__.flatMap(
						__.outE("tsw").filter(__.as("edge").select(last, "e").where(P.eq("edge")).by("speed")).
								group().by(__.inV()).
								by(__.project("curr", "time").by().
										by(__.sack(Operator.assign).by("depTime").select(last, "e").sack(Operator.minus).by("arrTime").sack()).
										filter(__.select("time").is(P.gte(0))).fold()).
								select(Column.values).unfold().order(local).by(__.select("time")).limit(local, 1).select("curr")
				).as("e").inV().simplePath()
		).times(20).map(__.union((Traversal) __.select(last, "e").by("speed"), (Traversal) __.path().by(T.id).by(__.valueMap("arrTime", "depTime"))).fold());
	}
}