Neo4j export

app/com/lynxanalytics/biggraph/frontend_operations/ExportOperations.scala

@@ -2,6 +2,7 @@ package com.lynxanalytics.biggraph.frontend_operations
 
 import com.lynxanalytics.biggraph.SparkFreeEnvironment
 import com.lynxanalytics.biggraph.graph_api.Scripting._
+import com.lynxanalytics.biggraph.graph_util.Scripting._
 import com.lynxanalytics.biggraph.graph_operations
 import com.lynxanalytics.biggraph.controllers._
 
@@ -153,4 +154,96 @@ class ExportOperations(env: SparkFreeEnvironment) extends OperationRegistry {
         .map(_ => wc.createSnapshotFromState(user, params("path"), getState))
     }
   })
+
+  abstract class Neo4jAttributeExport(context: Operation.Context) extends TriggerableOperation(context) {
+    params ++= List(
+      Param("url", "Neo4j connection", defaultValue = "bolt://localhost:7687"),
+      Param("username", "Neo4j username", defaultValue = "neo4j"),
+      Param("password", "Neo4j password", defaultValue = "neo4j"),
+      NonNegInt("version", "Snapshot revision ID", default = 1))
+    lazy val project = projectInput("graph")
+    val nodesOrRelationships: String
+    override def enabled = FEStatus.enabled
+    override def trigger(wc: WorkspaceController, gdc: GraphDrawingController) = {
+      gdc.getComputeBoxResult(List(exportResult.gUID))
+    }
+    override def getOutputs(): Map[BoxOutput, BoxOutputState] = {
+      params.validate()
+      Map(context.box.output(
+        context.meta.outputs(0)) -> BoxOutputState.from(exportResult, params.toMap - "password"))
+    }
+    def exportResult() = {
+      val keys = splitParam("keys")
+      val attrs = (splitParam("to_export") ++ keys).toSet.toList
+      val t = graph_operations.AttributesToTable.run(
+        attrs.map(a => a -> project.vertexAttributes(a)))
+      assert(keys.nonEmpty, "You have to choose one or more attributes to use as the keys for identifying the nodes in Neo4j.")
+      val op = graph_operations.ExportAttributesToNeo4j(
+        params("url"), params("username"), params("password"), params("labels"), keys,
+        params("version").toInt, nodesOrRelationships)
+      op(op.t, t).result.exportResult
+    }
+  }
+
+  registerOp(
+    "Export vertex attributes to Neo4j", defaultIcon, ExportOperations,
+    List("graph"), List("exported"), new Neo4jAttributeExport(_) {
+      import Operation.Implicits._
+      params ++= List(
+        Param("labels", "Node labels", defaultValue = ""),
+        Choice("keys", "Attribute to use as key",
+          // Cannot join on internal ID ("<id>") and stuff like that.
+          options = project.vertexAttrList.filter(!_.id.startsWith("<")), multipleChoice = true),
+        Choice("to_export", "Exported attributes", options = project.vertexAttrList, multipleChoice = true))
+      val nodesOrRelationships = "nodes"
+    })
+  registerOp(
+    "Export edge attributes to Neo4j", defaultIcon, ExportOperations,
+    List("graph"), List("exported"), new Neo4jAttributeExport(_) {
+      import Operation.Implicits._
+      params ++= List(
+        Param("labels", "Relationship labels", defaultValue = ""),
+        Choice("keys", "Attribute to use as key",
+          // Cannot join on internal ID ("<id>") and stuff like that.
+          options = project.edgeAttrList.filter(!_.id.startsWith("<")), multipleChoice = true),
+        Choice("to_export", "Exported attributes", options = project.edgeAttrList, multipleChoice = true))
+      val nodesOrRelationships = "edges"
+    })
+
+  registerOp(
+    "Export graph to Neo4j", defaultIcon, ExportOperations,
+    List("graph"), List("exported"), new TriggerableOperation(_) {
+      import Operation.Implicits._
+      params ++= List(
+        Param("url", "Neo4j connection", defaultValue = "bolt://localhost:7687"),
+        Param("username", "Neo4j username", defaultValue = "neo4j"),
+        Param("password", "Neo4j password", defaultValue = "neo4j"),
+        NonNegInt("version", "Snapshot revision ID", default = 1),
+        Choice("node_labels", "Attribute with node labels",
+          options = List(FEOption("", "")) ++ project.vertexAttrList[String]),
+        Choice("relationship_type", "Attribute with relationship type",
+          options = List(FEOption("", "")) ++ project.edgeAttrList[String]))
+      lazy val project = projectInput("graph")
+      override def enabled = FEStatus.enabled
+      override def trigger(wc: WorkspaceController, gdc: GraphDrawingController) = {
+        gdc.getComputeBoxResult(List(exportResult.gUID))
+      }
+      override def getOutputs(): Map[BoxOutput, BoxOutputState] = {
+        params.validate()
+        Map(context.box.output(
+          context.meta.outputs(0)) -> BoxOutputState.from(exportResult, params.toMap - "password"))
+      }
+      def exportResult() = {
+        val vsAttr = project.vertexAttributes.toMap +
+          (graph_operations.ExportGraphToNeo4j.VID -> project.vertexSet.idAttribute)
+        val esAttr = project.edgeAttributes.toMap +
+          (graph_operations.ExportGraphToNeo4j.SRCDST -> project.edgeBundle.srcDstAttribute)
+        val vs = graph_operations.AttributesToTable.run(vsAttr)
+        val es = graph_operations.AttributesToTable.run(esAttr)
+        val op = graph_operations.ExportGraphToNeo4j(
+          params("url"), params("username"), params("password"), params("node_labels"),
+          params("relationship_type"), params("version").toInt)
+        op(op.vs, vs)(op.es, es).result.exportResult
+      }
+    })
 }

app/com/lynxanalytics/biggraph/graph_operations/ExportToNeo4j.scala

@@ -0,0 +1,135 @@
+package com.lynxanalytics.biggraph.graph_operations
+
+import com.lynxanalytics.biggraph.graph_api._
+import com.lynxanalytics.biggraph.graph_util.HadoopFile
+import org.apache.spark
+import org.apache.spark.sql.SaveMode
+
+object ExportAttributesToNeo4j extends OpFromJson {
+  class Input extends MagicInputSignature {
+    val t = table
+  }
+  class Output(implicit instance: MetaGraphOperationInstance) extends MagicOutput(instance) {
+    val exportResult = scalar[String]
+  }
+  def fromJson(j: JsValue) = ExportAttributesToNeo4j(
+    (j \ "url").as[String], (j \ "username").as[String], (j \ "password").as[String],
+    (j \ "labels").as[String], (j \ "keys").as[Seq[String]], (j \ "version").as[Long],
+    (j \ "nodesOrRelationships").as[String])
+}
+
+case class Neo4jConnection(url: String, username: String, password: String) {
+  def send(df: spark.sql.DataFrame, query: String) {
+    df.write
+      .format("org.neo4j.spark.DataSource")
+      .option("authentication.type", "basic")
+      .option("authentication.basic.username", username)
+      .option("authentication.basic.password", password)
+      .option("url", url)
+      .option("query", query)
+      .save()
+  }
+}
+
+case class ExportAttributesToNeo4j(
+    url: String, username: String, password: String, labels: String, keys: Seq[String],
+    version: Long, nodesOrRelationships: String)
+  extends SparkOperation[ExportAttributesToNeo4j.Input, ExportAttributesToNeo4j.Output] {
+  val neo = Neo4jConnection(url, username, password)
+  @transient override lazy val inputs = new ExportAttributesToNeo4j.Input()
+  def outputMeta(instance: MetaGraphOperationInstance) = new ExportAttributesToNeo4j.Output()(instance)
+  override def toJson = Json.obj(
+    "url" -> url, "username" -> username, "password" -> password, "labels" -> labels,
+    "keys" -> keys, "version" -> version, "nodesOrRelationships" -> nodesOrRelationships)
+  def execute(
+    inputDatas: DataSet,
+    o: ExportAttributesToNeo4j.Output,
+    output: OutputBuilder,
+    rc: RuntimeContext): Unit = {
+    implicit val ds = inputDatas
+    // Drop null keys.
+    val df = keys.foldLeft(inputs.t.df)((df, key) => df.filter(df(key).isNotNull))
+    val keyMatch = keys.map(k => s"`$k`: event.`$k`").mkString(", ")
+    val query = nodesOrRelationships match {
+      case "nodes" => s"MATCH (n$labels {$keyMatch}) SET n += event"
+      case "relationships" => s"MATCH ()-[r$labels {$keyMatch}]-() SET r += event"
+    }
+    neo.send(df, query)
+    val exportResult = "Export done."
+    output(o.exportResult, exportResult)
+  }
+}
+
+object ExportGraphToNeo4j extends OpFromJson {
+  class Input extends MagicInputSignature {
+    val vs = table
+    val es = table
+  }
+  class Output(implicit instance: MetaGraphOperationInstance) extends MagicOutput(instance) {
+    val exportResult = scalar[String]
+  }
+  def fromJson(j: JsValue) = ExportGraphToNeo4j(
+    (j \ "url").as[String], (j \ "username").as[String], (j \ "password").as[String],
+    (j \ "nodeLabelsColumn").as[String], (j \ "relationshipTypeColumn").as[String],
+    (j \ "version").as[Long])
+  val VID = "!LynxKite ID"
+  val SRCDST = "!LynxKite endpoint IDs"
+  val SRCID = "!Source LynxKite ID"
+  val DSTID = "!Destination LynxKite ID"
+}
+
+case class ExportGraphToNeo4j(
+    url: String, username: String, password: String, nodeLabelsColumn: String,
+    relationshipTypeColumn: String, version: Long)
+  extends SparkOperation[ExportGraphToNeo4j.Input, ExportGraphToNeo4j.Output] {
+  import ExportGraphToNeo4j._
+  val neo = Neo4jConnection(url, username, password)
+  @transient override lazy val inputs = new ExportGraphToNeo4j.Input()
+  def outputMeta(instance: MetaGraphOperationInstance) = new ExportGraphToNeo4j.Output()(instance)
+  override def toJson = Json.obj(
+    "url" -> url, "username" -> username, "password" -> password,
+    "nodeLabelsColumn" -> nodeLabelsColumn, "relationshipTypeColumn" -> relationshipTypeColumn,
+    "version" -> version)
+  def execute(
+    inputDatas: DataSet,
+    o: ExportGraphToNeo4j.Output,
+    output: OutputBuilder,
+    rc: RuntimeContext): Unit = {
+    implicit val ds = inputDatas
+    val F = spark.sql.functions
+    // Prefix the internal IDs with this operation's GUID so different exports don't collide.
+    val guid = F.lit(this.gUID.toString + "-")
+    val vs = inputs.vs.df
+      .withColumn(VID, F.concat(guid, F.col(VID)))
+    val es = inputs.es.df
+      .withColumn(SRCID, F.concat(guid, F.col(SRCDST + "._1")))
+      .withColumn(DSTID, F.concat(guid, F.col(SRCDST + "._2")))
+      .drop(SRCDST)
+    if (nodeLabelsColumn.isEmpty) {
+      neo.send(vs, s"""
+        CREATE (n)
+        SET n += event
+      """)
+    } else {
+      neo.send(vs, s"""
+        CALL apoc.create.node(split(event.`$nodeLabelsColumn`, ','), event) YIELD node
+        RETURN 1
+      """)
+    }
+    if (relationshipTypeColumn.isEmpty) {
+      neo.send(es, s"""
+      MATCH (src {`$VID`: event.`$SRCID`}), (dst {`$VID`: event.`$DSTID`})
+      CREATE (src)-[r:EDGE]->(dst)
+      SET r += event
+    """)
+    } else {
+      neo.send(es, s"""
+      MATCH (src {`$VID`: event.`$SRCID`}), (dst {`$VID`: event.`$DSTID`})
+      CALL apoc.create.relationship(src, event.`$relationshipTypeColumn`, event, dst) YIELD rel
+      RETURN 1
+    """)
+    }
+    val exportResult = "Export done."
+    output(o.exportResult, exportResult)
+  }
+}

app/com/lynxanalytics/biggraph/graph_util/Scripting.scala

@@ -72,6 +72,11 @@ object Scripting {
       val op = graph_operations.AddReversedEdges()
       op(op.es, self).result.esPlus
     }
+
+    def srcDstAttribute: Attribute[(ID, ID)] = {
+      val op = graph_operations.EdgeBundleAsAttribute()
+      op(op.edges, self).result.attr
+    }
   }
 
   implicit class RichContainedAttribute[T](