Extra configuration for bigtable client

pyalex · pyalex · commit 0df33d010e00 · 2021-06-10T13:35:04.000+08:00
Signed-off-by: Oleksii Moskalenko &lt;moskalenko.alexey@gmail.com&gt;
diff --git a/spark/ingestion/src/main/scala/feast/ingestion/stores/bigtable/DefaultSource.scala b/spark/ingestion/src/main/scala/feast/ingestion/stores/bigtable/DefaultSource.scala
@@ -20,26 +20,38 @@ import org.apache.spark.sql.{DataFrame, SQLContext, SaveMode}
 import org.apache.spark.sql.sources.{BaseRelation, CreatableRelationProvider}
 import com.google.cloud.bigtable.hbase.BigtableConfiguration
 import feast.ingestion.stores.serialization.AvroSerializer
+import com.google.cloud.bigtable.hbase.BigtableOptionsFactory.{
+  BIGTABLE_BUFFERED_MUTATOR_ENABLE_THROTTLING,
+  BIGTABLE_BUFFERED_MUTATOR_THROTTLING_THRESHOLD_MILLIS,
+  BIGTABLE_BULK_MAX_ROW_KEY_COUNT,
+  BIGTABLE_DATA_CHANNEL_COUNT_KEY,
+  BIGTABLE_EMULATOR_HOST_KEY
+}
+import org.apache.hadoop.conf.Configuration
 
 class DefaultSource extends CreatableRelationProvider {
+  import DefaultSource._
+
   override def createRelation(
       sqlContext: SQLContext,
       mode: SaveMode,
       parameters: Map[String, String],
       data: DataFrame
   ): BaseRelation = {
     val bigtableConf = BigtableConfiguration.configure(
-      sqlContext.getConf("spark.bigtable.projectId"),
-      sqlContext.getConf("spark.bigtable.instanceId")
+      sqlContext.getConf(PROJECT_KEY),
+      sqlContext.getConf(INSTANCE_KEY)
     )
 
     if (sqlContext.getConf("spark.bigtable.emulatorHost", "").nonEmpty) {
       bigtableConf.set(
-        "google.bigtable.emulator.endpoint.host",
+        BIGTABLE_EMULATOR_HOST_KEY,
         sqlContext.getConf("spark.bigtable.emulatorHost")
       )
     }
 
+    configureBigTableClient(bigtableConf, sqlContext)
+
     val rel =
       new BigTableSinkRelation(
         sqlContext,
@@ -52,4 +64,32 @@ class DefaultSource extends CreatableRelationProvider {
     rel.insert(data, overwrite = false)
     rel
   }
+
+  private def configureBigTableClient(bigtableConf: Configuration, sqlContext: SQLContext): Unit = {
+    val confs = sqlContext.getAllConfs
+
+    confs.get(CHANNEL_COUNT_KEY).foreach(bigtableConf.set(BIGTABLE_DATA_CHANNEL_COUNT_KEY, _))
+    confs.get(MAX_ROW_COUNT_KEY).foreach(bigtableConf.set(BIGTABLE_BULK_MAX_ROW_KEY_COUNT, _))
+
+    confs
+      .get(ENABLE_THROTTLING_KEY)
+      .foreach(
+        bigtableConf.set(BIGTABLE_BUFFERED_MUTATOR_ENABLE_THROTTLING, _)
+      )
+    confs
+      .get(THROTTLING_THRESHOLD_MILLIS_KEY)
+      .foreach(
+        bigtableConf.set(BIGTABLE_BUFFERED_MUTATOR_THROTTLING_THRESHOLD_MILLIS, _)
+      )
+  }
+}
+
+object DefaultSource {
+  private val PROJECT_KEY  = "spark.bigtable.projectId"
+  private val INSTANCE_KEY = "spark.bigtable.instanceId"
+
+  private val CHANNEL_COUNT_KEY               = "spark.bigtable.channelCount"
+  private val ENABLE_THROTTLING_KEY           = "spark.bigtable.enableThrottling"
+  private val THROTTLING_THRESHOLD_MILLIS_KEY = "spark.bigtable.throttlingThresholdMs"
+  private val MAX_ROW_COUNT_KEY               = "spark.bigtable.maxRowCount"
 }
