Skip to content

Commit ac445cc

Browse files
authored
cherry pick support write back to original tag for jar (#34)
1 parent 0dc354c commit ac445cc

File tree

3 files changed

+20
-5
lines changed

3 files changed

+20
-5
lines changed

nebula-algorithm/src/main/resources/application.conf

+2
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,8 @@
4545
space:nb
4646
# Nebula tag name, the algorithm result will be write into this tag
4747
tag:pagerank
48+
# algorithm result is insert into new tag or update to original tag. type: insert/update
49+
type:insert
4850
}
4951
}
5052

nebula-algorithm/src/main/scala/com/vesoft/nebula/algorithm/config/Configs.scala

+11-3
Original file line numberDiff line numberDiff line change
@@ -99,8 +99,15 @@ object NebulaConfigEntry {
9999
val pswd = nebulaConfig.getString("write.pswd")
100100
val writeSpace = nebulaConfig.getString("write.space")
101101
val writeTag = nebulaConfig.getString("write.tag")
102+
val writeType = nebulaConfig.getString("write.type")
102103
val writeConfigEntry =
103-
NebulaWriteConfigEntry(graphAddress, writeMetaAddress, user, pswd, writeSpace, writeTag)
104+
NebulaWriteConfigEntry(graphAddress,
105+
writeMetaAddress,
106+
user,
107+
pswd,
108+
writeSpace,
109+
writeTag,
110+
writeType)
104111
NebulaConfigEntry(readConfigEntry, writeConfigEntry)
105112
}
106113
}
@@ -211,10 +218,11 @@ case class NebulaWriteConfigEntry(graphAddress: String = "",
211218
user: String = "",
212219
pswd: String = "",
213220
space: String = "",
214-
tag: String = "") {
221+
tag: String = "",
222+
writeType: String = "insert") {
215223
override def toString: String = {
216224
s"NebulaWriteConfigEntry: " +
217-
s"{graphAddress: $graphAddress, user: $user, password: $pswd, space: $space, tag: $tag}"
225+
s"{graphAddress: $graphAddress, user: $user, password: $pswd, space: $space, tag: $tag, type: $writeType}"
218226
}
219227
}
220228

nebula-algorithm/src/main/scala/com/vesoft/nebula/algorithm/writer/AlgoWriter.scala

+7-2
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
package com.vesoft.nebula.algorithm.writer
77

88
import com.vesoft.nebula.connector.connector.NebulaDataFrameWriter
9-
import com.vesoft.nebula.connector.{NebulaConnectionConfig, WriteNebulaVertexConfig}
9+
import com.vesoft.nebula.connector.{NebulaConnectionConfig, WriteMode, WriteNebulaVertexConfig}
1010
import com.vesoft.nebula.algorithm.config.{AlgoConstants, Configs}
1111
import org.apache.spark.sql.DataFrame
1212

@@ -22,6 +22,8 @@ class NebulaWriter(data: DataFrame, configs: Configs) extends AlgoWriter(data, c
2222
val tag = configs.nebulaConfig.writeConfigEntry.tag
2323
val user = configs.nebulaConfig.writeConfigEntry.user
2424
val passwd = configs.nebulaConfig.writeConfigEntry.pswd
25+
val writeType = configs.nebulaConfig.writeConfigEntry.writeType
26+
val writeMode = if (writeType.equals("insert")) WriteMode.INSERT else WriteMode.UPDATE
2527

2628
val config =
2729
NebulaConnectionConfig
@@ -30,13 +32,16 @@ class NebulaWriter(data: DataFrame, configs: Configs) extends AlgoWriter(data, c
3032
.withGraphAddress(graphAddress)
3133
.withConenctionRetry(2)
3234
.build()
33-
val nebulaWriteVertexConfig: WriteNebulaVertexConfig = WriteNebulaVertexConfig
35+
val nebulaWriteVertexConfig = WriteNebulaVertexConfig
3436
.builder()
37+
.withUser(user)
38+
.withPasswd(passwd)
3539
.withSpace(space)
3640
.withTag(tag)
3741
.withVidField(AlgoConstants.ALGO_ID_COL)
3842
.withVidAsProp(false)
3943
.withBatch(1000)
44+
.withWriteMode(writeMode)
4045
.build()
4146
data.write.nebula(config, nebulaWriteVertexConfig).writeVertices()
4247
}

0 commit comments

Comments
 (0)