fix: missing operateType

vesoft-inc · Nov 23, 2023 · 618b530 · 618b530
1 parent 1e681db
commit 618b530
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 7 deletions.
diff --git a/nebula-spark-common/src/main/scala/com/vesoft/nebula/connector/NebulaConfig.scala b/nebula-spark-common/src/main/scala/com/vesoft/nebula/connector/NebulaConfig.scala
@@ -712,28 +712,32 @@ class ReadNebulaConfig extends Serializable {
   var getPartitionNum: Int        = _
   var getLimit: Int               = _
   var getNgql: String             = _
+  var getOperateType: String      = _
   // todo add filter
   def this(space: String,
            label: String,
            returnCols: List[String],
            noColumn: Boolean,
            partitionNum: Int,
-           limit: Int) = {
+           limit: Int,
+           operateType: String) = {
     this()
     this.getSpace = space
     this.getLabel = label
     this.getReturnCols = returnCols
     this.getNoColumn = noColumn
     this.getPartitionNum = partitionNum
     this.getLimit = limit
+    this.getOperateType = operateType
   }
 
   def this(space: String,
            label: String,
            returnCols: List[String],
            noColumn: Boolean,
            ngql: String,
-           limit: Int) = {
+           limit: Int,
+           operateType: String) = {
     this()
     this.getNgql = ngql
     this.getSpace = space
@@ -742,6 +746,7 @@ class ReadNebulaConfig extends Serializable {
     this.getNoColumn = noColumn
     this.getLimit = limit
     this.getPartitionNum = 1
+    this.getOperateType = operateType
   }
 }
 
@@ -817,9 +822,9 @@ object ReadNebulaConfig {
     def build(): ReadNebulaConfig = {
       check()
       if (ngql != null && !ngql.isEmpty) {
-        new ReadNebulaConfig(space, label, returnCols.toList, noColumn, ngql, limit)
+        new ReadNebulaConfig(space, label, returnCols.toList, noColumn, ngql, limit, operateType)
       } else {
-        new ReadNebulaConfig(space, label, returnCols.toList, noColumn, partitionNum, limit)
+        new ReadNebulaConfig(space, label, returnCols.toList, noColumn, partitionNum, limit, operateType)
       }
     }
 

diff --git a/nebula-spark-common/src/main/scala/com/vesoft/nebula/connector/NebulaOptions.scala b/nebula-spark-common/src/main/scala/com/vesoft/nebula/connector/NebulaOptions.scala
@@ -18,20 +18,22 @@ class NebulaOptions(@transient val parameters: CaseInsensitiveMap[String]) exten
   import NebulaOptions._
 
   def this(parameters: Map[String, String], operaType: OperaType.Value) =
-    this(CaseInsensitiveMap(parameters))
+    this(CaseInsensitiveMap(parameters ++ Map(NebulaOptions.OPERATE_TYPE -> operaType.toString)))
 
   def this(hostAndPorts: String,
            spaceName: String,
            dataType: String,
            label: String,
-           parameters: Map[String, String]) = {
+           parameters: Map[String, String],
+           operateType: OperaType.Value) = {
     this(
       CaseInsensitiveMap(
         parameters ++ Map(
           NebulaOptions.META_ADDRESS -> hostAndPorts,
           NebulaOptions.SPACE_NAME   -> spaceName,
           NebulaOptions.TYPE         -> dataType,
-          NebulaOptions.LABEL        -> label
+          NebulaOptions.LABEL        -> label,
+          NebulaOptions.OPERATE_TYPE -> operateType.toString
         ))
     )
   }