add multi-column tests

core/src/test/scala/org/apache/spark/sql/BaseTestGenerationSpec.scala

@@ -1,6 +1,6 @@
 /*
  *
- * Copyright 2017 PingCAP, Inc.
+ * Copyright 2019 PingCAP, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -21,12 +21,15 @@ trait BaseTestGenerationSpec {
 
   protected val rowCount: Int = 50
 
-  protected def getTableName(dataType: String): String = s"test_$dataType"
+  private def toString(dataTypes: Seq[String]) = dataTypes.map(_.head).mkString("_")
 
-  protected def getTableName(dataType: String, desc: String): String = s"test_${desc}_$dataType"
+  protected def getTableName(dataTypes: String*): String = s"test_${toString(dataTypes)}"
+
+  protected def getTableNameWithDesc(desc: String, dataTypes: String*): String =
+    s"test_${desc}_${toString(dataTypes)}"
 
   protected def getColumnName(dataType: String): String = s"col_$dataType"
 
-  protected def getIndexName(dataType: String): String = s"idx_$dataType"
+  protected def getIndexName(dataTypes: String*): String = s"idx_${toString(dataTypes)}"
 
 }

core/src/test/scala/org/apache/spark/sql/MysqlDataTypeSuite.scala

@@ -1,4 +1,5 @@
 /*
+ *
  * Copyright 2019 PingCAP, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
@@ -11,6 +12,7 @@
  * distributed under the License is distributed on an "AS IS" BASIS,
  * See the License for the specific language governing permissions and
  * limitations under the License.
+ *
  */
 
 package org.apache.spark.sql

core/src/test/scala/org/apache/spark/sql/TiSparkTestAction.scala

@@ -0,0 +1,29 @@
+/*
+ *
+ * Copyright 2019 PingCAP, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package org.apache.spark.sql
+
+import scala.util.Random
+
+trait TiSparkTestAction {
+  val database: String
+  val testDesc: String
+  // Randomizer for tests
+  val r: Random = new Random(1234)
+
+  def test(): Unit
+}

core/src/test/scala/org/apache/spark/sql/test/SharedSQLContext.scala

@@ -26,7 +26,6 @@ import com.pingcap.tispark.TiDBUtils
 import com.pingcap.tispark.statistics.StatisticsManager
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql._
-import org.apache.spark.sql.catalyst.util.resourceToString
 import org.apache.spark.sql.test.TestConstants._
 import org.apache.spark.sql.test.Utils._
 import org.apache.spark.{SparkConf, SparkContext, SparkFunSuite}

core/src/test/scala/org/apache/spark/sql/test/generator/Data.scala

@@ -1,4 +1,5 @@
 /*
+ *
  * Copyright 2019 PingCAP, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
@@ -11,6 +12,7 @@
  * distributed under the License is distributed on an "AS IS" BASIS,
  * See the License for the specific language governing permissions and
  * limitations under the License.
+ *
  */
 
 package org.apache.spark.sql.test.generator

core/src/test/scala/org/apache/spark/sql/test/generator/DataType.scala

@@ -1,4 +1,5 @@
 /*
+ *
  * Copyright 2019 PingCAP, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
@@ -11,6 +12,7 @@
  * distributed under the License is distributed on an "AS IS" BASIS,
  * See the License for the specific language governing permissions and
  * limitations under the License.
+ *
  */
 
 package org.apache.spark.sql.test.generator

core/src/test/scala/org/apache/spark/sql/test/generator/Index.scala

@@ -1,4 +1,5 @@
 /*
+ *
  * Copyright 2019 PingCAP, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
@@ -11,6 +12,7 @@
  * distributed under the License is distributed on an "AS IS" BASIS,
  * See the License for the specific language governing permissions and
  * limitations under the License.
+ *
  */
 
 package org.apache.spark.sql.test.generator

core/src/test/scala/org/apache/spark/sql/test/generator/IndexColumn.scala

@@ -1,4 +1,5 @@
 /*
+ *
  * Copyright 2019 PingCAP, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
@@ -11,6 +12,7 @@
  * distributed under the License is distributed on an "AS IS" BASIS,
  * See the License for the specific language governing permissions and
  * limitations under the License.
+ *
  */
 
 package org.apache.spark.sql.test.generator
@@ -54,14 +56,15 @@ case class ColumnInfo(columnName: String,
     }
   }
 
-  private val baseType = getBaseType(dataType)
-
-  private val (len, decimal): (Long, Int) = if (length._1 == null) {
-    (getLength(baseType), getDecimal(baseType))
-  } else if (length._2 == null) {
-    (length._1.toLong, getDecimal(baseType))
-  } else {
-    (length._1.toLong, length._2)
+  val (len, decimal): (Long, Int) = {
+    val baseType = getBaseType(dataType)
+    if (length._1 == null) {
+      (getLength(baseType), getDecimal(baseType))
+    } else if (length._2 == null) {
+      (length._1.toLong, getDecimal(baseType))
+    } else {
+      (length._1.toLong, length._2)
+    }
   }
 
   {

core/src/test/scala/org/apache/spark/sql/test/generator/Schema.scala

@@ -1,4 +1,5 @@
 /*
+ *
  * Copyright 2019 PingCAP, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
@@ -11,6 +12,7 @@
  * distributed under the License is distributed on an "AS IS" BASIS,
  * See the License for the specific language governing permissions and
  * limitations under the License.
+ *
  */
 
 package org.apache.spark.sql.test.generator

core/src/test/scala/org/apache/spark/sql/test/generator/TestDataGenerator.scala

@@ -1,4 +1,5 @@
 /*
+ *
  * Copyright 2019 PingCAP, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
@@ -11,6 +12,7 @@
  * distributed under the License is distributed on an "AS IS" BASIS,
  * See the License for the specific language governing permissions and
  * limitations under the License.
+ *
  */
 
 package org.apache.spark.sql.test.generator
@@ -268,29 +270,82 @@ object TestDataGenerator {
                                   offset: Int,
                                   r: Random,
                                   valueGenerator: ValueGenerator): Unit = {
-    if (valueGenerator.isPrimaryKey) {
-      assert(!valueGenerator.nullable, "Generate fails: Value cannot be null for primary key")
-      val value = valueGenerator.randomUniqueValue(r)
-      row.set(offset, valueGenerator.tiDataType, value)
+    val value = valueGenerator.next(r)
+    if (value == null) {
+      row.setNull(offset)
     } else {
-      if (valueGenerator.randomNull(r)) {
-        row.setNull(offset)
-      } else {
-        val value = valueGenerator.randomValue(r)
-        row.set(offset, valueGenerator.tiDataType, value)
+      row.set(offset, valueGenerator.tiDataType, value)
+    }
+  }
+
+  def hash(value: Any): String = value match {
+    case null           => "null"
+    case b: Array[Byte] => b.mkString("[", ",", "]")
+    case list: List[Any] =>
+      val ret = StringBuilder.newBuilder
+      ret ++= "("
+      for (i <- list.indices) {
+        if (i > 0) ret ++= ","
+        ret ++= hash(list(i))
       }
+      ret ++= ")"
+      ret.toString
+    case x => x.toString
+  }
+
+  def checkUnique(value: Any, set: mutable.Set[Any]): Boolean = {
+    val hashedValue = hash(value)
+    if (!set.apply(hashedValue)) {
+      set += hashedValue
+      true
+    } else {
+      false
     }
   }
 
-  private def generateRandomRows(schema: Schema, n: Long, r: Random): List[TiRow] = {
-    (1.toLong to n).map { _ =>
-      val length = schema.columnInfo.length
-      val row: TiRow = ObjectRowImpl.create(length)
+  private def generateRandomRow(schema: Schema,
+                                r: Random,
+                                pkOffset: List[Int],
+                                set: mutable.Set[Any]): TiRow = {
+    val length = schema.columnInfo.length
+    val row: TiRow = ObjectRowImpl.create(length)
+    while (true) {
       for (i <- schema.columnInfo.indices) {
         val columnInfo = schema.columnInfo(i)
         generateRandomValue(row, i, r, columnInfo.generator)
       }
-      row
+      if (pkOffset.nonEmpty) {
+        val value = pkOffset.map { i =>
+          row.get(i, schema.columnInfo(i).generator.tiDataType)
+        }
+        if (checkUnique(value, set)) {
+          return row
+        }
+      } else {
+        return row
+      }
+    }
+    throw new RuntimeException("Inaccessible")
+  }
+
+  private def generateRandomRows(schema: Schema, n: Long, r: Random): List[TiRow] = {
+    val set: mutable.Set[Any] = mutable.HashSet.empty[Any]
+    // offset of pk columns
+    val pkOffset: List[Int] = {
+      val primary = schema.indexInfo.filter(_.isPrimary)
+      if (primary.nonEmpty && primary.size == 1) {
+        primary.head.indexColumns.map(x => schema.columnNames.indexOf(x.column))
+      } else {
+        List.empty[Int]
+      }
+    }
+    schema.columnInfo.foreach { col =>
+      col.generator.reset()
+      col.generator.preGenerateRandomValues(r, n)
+    }
+
+    (1.toLong to n).map { _ =>
+      generateRandomRow(schema, r, pkOffset, set)
     }.toList
   }