add more test for utils

nebula-spark-connector/src/test/resources/edge.csv

@@ -0,0 +1,14 @@
+id1,id2,col1,col2,col3,col4,col5,col6,col7,col8,col9,col10,col11,col12,col13
+1,2,Tom,tom,10,20,30,40,2021-01-27,2021-01-01T12:10:10,43535232,true,1.0,2.0,10:10:10
+2,3,Jina,Jina,11,21,31,41,2021-01-28,2021-01-02T12:10:10,43535232,false,1.1,2.1,11:10:10
+3,4,Tim,Tim,12,22,32,42,2021-01-29,2021-01-03T12:10:10,43535232,false,1.2,2.2,12:10:10
+4,5,张三,张三,13,23,33,43,2021-01-30,2021-01-04T12:10:10,43535232,true,1.3,2.3,13:10:10
+5,6,李四,李四,14,24,34,44,2021-02-01,2021-01-05T12:10:10,43535232,false,1.4,2.4,14:10:10
+6,7,王五,王五,15,25,35,45,2021-02-02,2021-01-06T12:10:10,0,false,1.5,2.5,15:10:10
+7,1,Jina,Jina,16,26,36,46,2021-02-03,2021-01-07T12:10:10,43535232,true,1.6,2.6,16:10:10
+8,1,Jina,Jina,17,27,37,47,2021-02-04,2021-01-08T12:10:10,43535232,false,1.7,2.7,17:10:10
+9,1,Jina,Jina,18,28,38,48,2021-02-05,2021-01-09T12:10:10,43535232,true,1.8,2.8,18:10:10
+10,2,Jina,Jina,19,29,39,49,2021-02-06,2021-01-10T12:10:10,43535232,false,1.9,2.9,19:10:10
+-1,5,Jina,Jina,20,30,40,50,2021-02-07,2021-02-11T12:10:10,43535232,false,2.0,3.0,20:10:10
+-2,6,Jina,Jina,21,31,41,51,2021-02-08,2021-03-12T12:10:10,43535232,false,2.1,3.1,21:10:10
+-3,7,Jina,Jina,22,32,42,52,2021-02-09,2021-04-13T12:10:10,43535232,false,2.2,3.2,22:10:10

nebula-spark-connector/src/test/resources/vertex.csv

@@ -0,0 +1,14 @@
+id,col1,col2,col3,col4,col5,col6,col7,col8,col9,col10,col11,col12,col13
+1,Tom,tom,10,20,30,40,2021-01-27,2021-01-01T12:10:10,43535232,true,1.0,2.0,10:10:10
+2,Jina,Jina,11,21,31,41,2021-01-28,2021-01-02T12:10:10,43535232,false,1.1,2.1,11:10:10
+3,Tim,Tim,12,22,32,42,2021-01-29,2021-01-03T12:10:10,43535232,false,1.2,2.2,12:10:10
+4,张三,张三,13,23,33,43,2021-01-30,2021-01-04T12:10:10,43535232,true,1.3,2.3,13:10:10
+5,李四,李四,14,24,34,44,2021-02-01,2021-01-05T12:10:10,43535232,false,1.4,2.4,14:10:10
+6,王五,王五,15,25,35,45,2021-02-02,2021-01-06T12:10:10,0,false,1.5,2.5,15:10:10
+7,Jina,Jina,16,26,36,46,2021-02-03,2021-01-07T12:10:10,43535232,true,1.6,2.6,16:10:10
+8,Jina,Jina,17,27,37,47,2021-02-04,2021-01-08T12:10:10,43535232,false,1.7,2.7,17:10:10
+9,Jina,Jina,18,28,38,48,2021-02-05,2021-01-09T12:10:10,43535232,true,1.8,2.8,18:10:10
+10,Jina,Jina,19,29,39,49,2021-02-06,2021-01-10T12:10:10,43535232,false,1.9,2.9,19:10:10
+-1,Jina,Jina,20,30,40,50,2021-02-07,2021-02-11T12:10:10,43535232,false,2.0,3.0,20:10:10
+-2,Jina,Jina,21,31,41,51,2021-02-08,2021-03-12T12:10:10,43535232,false,2.1,3.1,21:10:10
+-3,Jina,Jina,22,32,42,52,2021-02-09,2021-04-13T12:10:10,43535232,false,2.2,3.2,22:10:10

nebula-spark-connector/src/test/scala/com/vesoft/nebula/connector/DataTypeEnumSuite.scala

@@ -0,0 +1,21 @@
+/* Copyright (c) 2021 vesoft inc. All rights reserved.
+ *
+ * This source code is licensed under Apache 2.0 License,
+ * attached with Common Clause Condition 1.0, found in the LICENSES directory.
+ */
+
+package com.vesoft.nebula.connector
+
+import org.scalatest.funsuite.AnyFunSuite
+
+class DataTypeEnumSuite extends AnyFunSuite {
+
+  test("validDataType") {
+    assert(DataTypeEnum.validDataType("vertex"))
+    assert(DataTypeEnum.validDataType("VERTEX"))
+    assert(DataTypeEnum.validDataType("edge"))
+    assert(DataTypeEnum.validDataType("EDGE"))
+    assert(!DataTypeEnum.validDataType("relation"))
+  }
+
+}

nebula-spark-connector/src/test/scala/com/vesoft/nebula/connector/NebulaUtilsSuite.scala

@@ -0,0 +1,97 @@
+/* Copyright (c) 2021 vesoft inc. All rights reserved.
+ *
+ * This source code is licensed under Apache 2.0 License,
+ * attached with Common Clause Condition 1.0, found in the LICENSES directory.
+ */
+
+package com.vesoft.nebula.connector
+
+import com.vesoft.nebula.meta.{ColumnDef, ColumnTypeDef, PropertyType}
+import org.apache.spark.sql.types.{
+  BooleanType,
+  DoubleType,
+  LongType,
+  StringType,
+  StructField,
+  StructType
+}
+import org.scalatest.funsuite.AnyFunSuite
+
+class NebulaUtilsSuite extends AnyFunSuite {
+
+  test("convertDataType") {
+    assert(NebulaUtils.convertDataType(new ColumnTypeDef(PropertyType.VID)) == LongType)
+    assert(NebulaUtils.convertDataType(new ColumnTypeDef(PropertyType.INT8)) == LongType)
+    assert(NebulaUtils.convertDataType(new ColumnTypeDef(PropertyType.INT16)) == LongType)
+    assert(NebulaUtils.convertDataType(new ColumnTypeDef(PropertyType.INT32)) == LongType)
+    assert(NebulaUtils.convertDataType(new ColumnTypeDef(PropertyType.INT64)) == LongType)
+    assert(NebulaUtils.convertDataType(new ColumnTypeDef(PropertyType.TIMESTAMP)) == LongType)
+    assert(NebulaUtils.convertDataType(new ColumnTypeDef(PropertyType.BOOL)) == BooleanType)
+    assert(NebulaUtils.convertDataType(new ColumnTypeDef(PropertyType.FLOAT)) == DoubleType)
+    assert(NebulaUtils.convertDataType(new ColumnTypeDef(PropertyType.DOUBLE)) == DoubleType)
+    assert(NebulaUtils.convertDataType(new ColumnTypeDef(PropertyType.FIXED_STRING)) == StringType)
+    assert(NebulaUtils.convertDataType(new ColumnTypeDef(PropertyType.STRING)) == StringType)
+    assert(NebulaUtils.convertDataType(new ColumnTypeDef(PropertyType.DATE)) == StringType)
+    assert(NebulaUtils.convertDataType(new ColumnTypeDef(PropertyType.DATETIME)) == StringType)
+    assert(NebulaUtils.convertDataType(new ColumnTypeDef(PropertyType.TIME)) == StringType)
+    assertThrows[IllegalArgumentException](
+      NebulaUtils.convertDataType(new ColumnTypeDef(PropertyType.UNKNOWN)))
+  }
+
+  test("getColDataType") {
+    val columnDefs: List[ColumnDef] = List(
+      new ColumnDef("col1".getBytes(), new ColumnTypeDef(PropertyType.INT8)),
+      new ColumnDef("col2".getBytes(), new ColumnTypeDef(PropertyType.DOUBLE)),
+      new ColumnDef("col3".getBytes(), new ColumnTypeDef(PropertyType.STRING)),
+      new ColumnDef("col4".getBytes(), new ColumnTypeDef(PropertyType.DATE)),
+      new ColumnDef("col5".getBytes(), new ColumnTypeDef(PropertyType.DATETIME)),
+      new ColumnDef("col6".getBytes(), new ColumnTypeDef(PropertyType.TIME)),
+      new ColumnDef("col7".getBytes(), new ColumnTypeDef(PropertyType.TIMESTAMP)),
+      new ColumnDef("col8".getBytes(), new ColumnTypeDef(PropertyType.BOOL))
+    )
+    assert(NebulaUtils.getColDataType(columnDefs, "col1") == LongType)
+    assert(NebulaUtils.getColDataType(columnDefs, "col2") == DoubleType)
+    assert(NebulaUtils.getColDataType(columnDefs, "col3") == StringType)
+    assert(NebulaUtils.getColDataType(columnDefs, "col4") == StringType)
+    assert(NebulaUtils.getColDataType(columnDefs, "col5") == StringType)
+    assert(NebulaUtils.getColDataType(columnDefs, "col6") == StringType)
+    assert(NebulaUtils.getColDataType(columnDefs, "col7") == LongType)
+    assert(NebulaUtils.getColDataType(columnDefs, "col8") == BooleanType)
+    assertThrows[IllegalArgumentException](NebulaUtils.getColDataType(columnDefs, "col9"))
+  }
+
+  test("makeGetters") {
+    val schema = StructType(
+      List(
+        StructField("col1", LongType, nullable = false),
+        StructField("col2", LongType, nullable = true)
+      ))
+    assert(NebulaUtils.makeGetters(schema).length == 2)
+  }
+
+  test("isNumic") {
+    assert(NebulaUtils.isNumic("123"))
+    assert(NebulaUtils.isNumic("-123"))
+    assert(!NebulaUtils.isNumic("1.0"))
+    assert(!NebulaUtils.isNumic("a123"))
+    assert(!NebulaUtils.isNumic("123b"))
+  }
+
+  test("escapeUtil") {
+    assert(NebulaUtils.escapeUtil("123").equals("123"))
+    // a\bc -> a\\bc
+    assert(NebulaUtils.escapeUtil("a\bc").equals("a\\bc"))
+    // a\tbc -> a\\tbc
+    assert(NebulaUtils.escapeUtil("a\tbc").equals("a\\tbc"))
+    // a\nbc -> a\\nbc
+    assert(NebulaUtils.escapeUtil("a\nbc").equals("a\\nbc"))
+    // a\"bc -> a\\"bc
+    assert(NebulaUtils.escapeUtil("a\"bc").equals("a\\\"bc"))
+    // a\'bc -> a\\'bc
+    assert(NebulaUtils.escapeUtil("a\'bc").equals("a\\'bc"))
+    // a\rbc -> a\\rbc
+    assert(NebulaUtils.escapeUtil("a\rbc").equals("a\\rbc"))
+    // a\bbc -> a\\bbc
+    assert(NebulaUtils.escapeUtil("a\bbc").equals("a\\bbc"))
+  }
+}

nebula-spark-connector/src/test/scala/com/vesoft/nebula/connector/PartitionUtilsSuite.scala

@@ -0,0 +1,117 @@
+/* Copyright (c) 2021 vesoft inc. All rights reserved.
+ *
+ * This source code is licensed under Apache 2.0 License,
+ * attached with Common Clause Condition 1.0, found in the LICENSES directory.
+ */
+
+package com.vesoft.nebula.connector
+
+import org.scalatest.funsuite.AnyFunSuite
+
+/**
+  * base data: spark partition is 10
+  */
+class PartitionUtilsSuite extends AnyFunSuite {
+  val partition: Int = 10
+
+  test("getScanParts: nebula part is the same with spark partition") {
+    val nebulaPart: Int = 10
+    val partsForIndex0  = PartitionUtils.getScanParts(1, nebulaPart, partition)
+    assert(partsForIndex0.size == 1)
+    assert(partsForIndex0.head == 1)
+
+    val partsForIndex1 = PartitionUtils.getScanParts(2, nebulaPart, partition)
+    assert(partsForIndex1.head == 2)
+
+    val partsForIndex2 = PartitionUtils.getScanParts(3, nebulaPart, partition)
+    assert(partsForIndex2.head == 3)
+
+    val partsForIndex3 = PartitionUtils.getScanParts(4, nebulaPart, partition)
+    assert(partsForIndex3.head == 4)
+
+    val partsForIndex4 = PartitionUtils.getScanParts(5, nebulaPart, partition)
+    assert(partsForIndex4.head == 5)
+
+    val partsForIndex5 = PartitionUtils.getScanParts(6, nebulaPart, partition)
+    assert(partsForIndex5.head == 6)
+
+    val partsForIndex6 = PartitionUtils.getScanParts(7, nebulaPart, partition)
+    assert(partsForIndex6.head == 7)
+
+    val partsForIndex7 = PartitionUtils.getScanParts(8, nebulaPart, partition)
+    assert(partsForIndex7.head == 8)
+
+    val partsForIndex8 = PartitionUtils.getScanParts(9, nebulaPart, partition)
+    assert(partsForIndex8.head == 9)
+
+    val partsForIndex9 = PartitionUtils.getScanParts(10, nebulaPart, partition)
+    assert(partsForIndex9.head == 10)
+  }
+
+  test("getScanParts: nebula part is more than spark partition") {
+    val nebulaPart: Int = 20
+    val partsForIndex0  = PartitionUtils.getScanParts(1, nebulaPart, partition)
+    assert(partsForIndex0.contains(1) && partsForIndex0.contains(11))
+
+    val partsForIndex1 = PartitionUtils.getScanParts(2, nebulaPart, partition)
+    assert(partsForIndex1.contains(2) && partsForIndex1.contains(12))
+
+    val partsForIndex2 = PartitionUtils.getScanParts(3, nebulaPart, partition)
+    assert(partsForIndex2.contains(3) && partsForIndex2.contains(13))
+
+    val partsForIndex3 = PartitionUtils.getScanParts(4, nebulaPart, partition)
+    assert(partsForIndex3.contains(4) && partsForIndex3.contains(14))
+
+    val partsForIndex4 = PartitionUtils.getScanParts(5, nebulaPart, partition)
+    assert(partsForIndex4.contains(5) && partsForIndex4.contains(15))
+
+    val partsForIndex5 = PartitionUtils.getScanParts(6, nebulaPart, partition)
+    assert(partsForIndex5.contains(6) && partsForIndex5.contains(16))
+
+    val partsForIndex6 = PartitionUtils.getScanParts(7, nebulaPart, partition)
+    assert(partsForIndex6.contains(7) && partsForIndex6.contains(17))
+
+    val partsForIndex7 = PartitionUtils.getScanParts(8, nebulaPart, partition)
+    assert(partsForIndex7.contains(8) && partsForIndex7.contains(18))
+
+    val partsForIndex8 = PartitionUtils.getScanParts(9, nebulaPart, partition)
+    assert(partsForIndex8.contains(9) && partsForIndex8.contains(19))
+
+    val partsForIndex9 = PartitionUtils.getScanParts(10, nebulaPart, partition)
+    assert(partsForIndex9.contains(10) && partsForIndex9.contains(20))
+  }
+
+  test("getScanParts: nebula part is less than spark partition") {
+    val nebulaPart: Int = 5
+    val partsForIndex0  = PartitionUtils.getScanParts(1, nebulaPart, partition)
+    assert(partsForIndex0.contains(1))
+
+    val partsForIndex1 = PartitionUtils.getScanParts(2, nebulaPart, partition)
+    assert(partsForIndex1.contains(2))
+
+    val partsForIndex2 = PartitionUtils.getScanParts(3, nebulaPart, partition)
+    assert(partsForIndex2.contains(3))
+
+    val partsForIndex3 = PartitionUtils.getScanParts(4, nebulaPart, partition)
+    assert(partsForIndex3.contains(4))
+
+    val partsForIndex4 = PartitionUtils.getScanParts(5, nebulaPart, partition)
+    assert(partsForIndex4.contains(5))
+
+    val partsForIndex5 = PartitionUtils.getScanParts(6, nebulaPart, partition)
+    assert(partsForIndex5.isEmpty)
+
+    val partsForIndex6 = PartitionUtils.getScanParts(7, nebulaPart, partition)
+    assert(partsForIndex6.isEmpty)
+
+    val partsForIndex7 = PartitionUtils.getScanParts(8, nebulaPart, partition)
+    assert(partsForIndex7.isEmpty)
+
+    val partsForIndex8 = PartitionUtils.getScanParts(9, nebulaPart, partition)
+    assert(partsForIndex8.isEmpty)
+
+    val partsForIndex9 = PartitionUtils.getScanParts(10, nebulaPart, partition)
+    assert(partsForIndex9.isEmpty)
+  }
+
+}