support overriding druidRelation options via sqlContext variables

hbutani · Jun 3, 2016 · 01b42d4 · 01b42d4
1 parent 4294698
commit 01b42d4
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/src/main/scala/org/sparklinedata/druid/DruidRDD.scala b/src/main/scala/org/sparklinedata/druid/DruidRDD.scala
@@ -72,6 +72,7 @@ class DruidRDD(sqlContext: SQLContext,
                 val dQuery : DruidQuery)  extends  RDD[InternalRow](sqlContext.sparkContext, Nil) {
 
   val druidQueryAcc : DruidQueryExecutionMetric = new DruidQueryExecutionMetric()
+  val numSegmentsPerQuery = drInfo.options.numSegmentsPerHistoricalQuery(sqlContext)
 
   @DeveloperApi
   override def compute(split: Partition, context: TaskContext): Iterator[InternalRow] = {
@@ -129,7 +130,6 @@ class DruidRDD(sqlContext: SQLContext,
       dQuery.intervalSplits
     )
       var idx = -1
-      val numSegmentsPerQuery = drInfo.options.numSegmentsPerHistoricalQuery
 
       val l  = (for(
         hA <- hAssigns;

diff --git a/src/main/scala/org/sparklinedata/druid/metadata/DruidRelationInfo.scala b/src/main/scala/org/sparklinedata/druid/metadata/DruidRelationInfo.scala
@@ -43,6 +43,7 @@ case class DruidRelationInfo(val fullName : DruidRelationName,
 
   def sourceDF(sqlContext : SQLContext) = sqlContext.table(sourceDFName)
 
+
 }
 
 case class DruidRelationOptions(val maxCardinality : Long,
@@ -55,7 +56,18 @@ case class DruidRelationOptions(val maxCardinality : Long,
                                 zkDruidPath : String,
                                 queryHistoricalServers : Boolean,
                                 zkQualifyDiscoveryNames : Boolean,
-                                numSegmentsPerHistoricalQuery : Int)
+                                numSegmentsPerHistoricalQuery : Int) {
+
+  def sqlContextOption(nm : String) = s"spark.sparklinedata.druid.option.$nm"
+
+  def numSegmentsPerHistoricalQuery(sqlContext : SQLContext) : Int = {
+    sqlContext.getConf(
+      sqlContextOption("numSegmentsPerHistoricalQuery"),
+        numSegmentsPerHistoricalQuery.toString
+    ).toInt
+  }
+
+}
 
 private[druid] object MappingBuilder extends Logging {