Add debug utils (#136)

pingcap · Dec 21, 2017 · b23ea86 · b23ea86
1 parent c65a12e
commit b23ea86
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/src/main/scala/org/apache/spark/sql/TiContext.scala b/src/main/scala/org/apache/spark/sql/TiContext.scala
@@ -15,21 +15,35 @@
 
 package org.apache.spark.sql
 
+
+import com.pingcap.tikv.tools.RegionUtils
+
 import com.pingcap.tikv.{TiConfiguration, TiSession}
 import com.pingcap.tispark._
 import org.apache.spark.SparkConf
 import org.apache.spark.internal.Logging
+import scala.collection.JavaConverters._
 
 class TiContext(val session: SparkSession) extends Serializable with Logging {
   val sqlContext: SQLContext = session.sqlContext
   val conf: SparkConf = session.sparkContext.conf
   val tiConf: TiConfiguration = TiUtils.sparkConfToTiConf(conf)
   val tiSession: TiSession = TiSession.create(tiConf)
   val meta: MetaManager = new MetaManager(tiSession.getCatalog)
+
+  val debug: DebugTool = new DebugTool
+
   TiUtils.sessionInitialize(session, tiSession)
 
+
   final val version: String = TiSparkVersion.version
 
+  class DebugTool {
+    def getRegionDistribution(dbName: String, tableName: String): Map[String, Integer] = {
+      RegionUtils.getRegionDistribution(tiSession, dbName, tableName).asScala.toMap
+    }
+  }
+
   def tidbTable(dbName: String, tableName: String): DataFrame = {
     val tiRelation = new TiDBRelation(
       tiSession,