add default max batch size and batchSize parameter for SF updates

README.md

@@ -62,6 +62,8 @@ $ bin/spark-shell --packages com.springml:spark-salesforce_2.11:1.1.3
 * `timeout`: (Optional) The maximum time spent polling for the completion of bulk query job. This option can only be used when `bulk` is `true`.
 * `externalIdFieldName`: (Optional) The name of the field used as the external ID for Salesforce Object. This value is only used when doing an update or upsert. Default "Id".
 * `queryAll`: (Optional) Toggle to retrieve deleted and archived records for SOQL queries. Default value is `false`.
+### Options only supported for fetching Salesforce Objects.
+* `batchSize`: (Optional) maximum number of records per batch when performing updates. Defaults to 5000 (note that batches greater than 10000 will result in a error)
 
 
 ### Scala API

build.sbt

@@ -7,11 +7,12 @@ organization := "com.springml"
 scalaVersion := "2.11.8"
 
 resolvers += "sonatype-snapshots" at "https://oss.sonatype.org/content/repositories/snapshots/"
+resolvers += "jitpack" at "https://jitpack.io"
 
 libraryDependencies ++= Seq(
   "com.force.api" % "force-wsc" % "40.0.0",
   "com.force.api" % "force-partner-api" % "40.0.0",
-  "com.springml" % "salesforce-wave-api" % "1.0.10",
+  "com.github.loanpal-engineering" % "salesforce-wave-api" % "eb71436",
   "org.mockito" % "mockito-core" % "2.0.31-beta"
 )
 

src/main/scala/com/springml/spark/salesforce/DatasetRelation.scala

@@ -116,7 +116,9 @@ case class DatasetRelation(
 
   private def cast(fieldValue: String, toType: DataType,
       nullable: Boolean = true, fieldName: String): Any = {
-    if (fieldValue == "" && nullable && !toType.isInstanceOf[StringType]) {
+    if (fieldValue == null)
+      null
+    else if (fieldValue == "" && nullable && !toType.isInstanceOf[StringType]) {
       null
     } else {
       toType match {

src/main/scala/com/springml/spark/salesforce/DefaultSource.scala

@@ -26,6 +26,7 @@ import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.{DataFrame, SQLContext, SaveMode}
 
 import scala.collection.mutable.ListBuffer
+import scala.util.{Failure, Success, Try}
 
 /**
  * Default source for Salesforce wave data source.
@@ -123,6 +124,15 @@ class DefaultSource extends RelationProvider with SchemaRelationProvider with Cr
     val encodeFields = parameters.get("encodeFields")
     val monitorJob = parameters.getOrElse("monitorJob", "false")
     val externalIdFieldName = parameters.getOrElse("externalIdFieldName", "Id")
+    val batchSizeStr = parameters.getOrElse("batchSize", "5000")
+    val batchSize = Try(batchSizeStr.toInt) match {
+      case Success(v)=> v
+      case Failure(e)=> {
+        val errorMsg = "batchSize parameter not an integer."
+        logger.error(errorMsg)
+        throw new Exception(errorMsg)
+      }
+    }
 
     validateMutualExclusive(datasetName, sfObject, "datasetName", "sfObject")
 
@@ -141,7 +151,7 @@ class DefaultSource extends RelationProvider with SchemaRelationProvider with Cr
     } else {
       logger.info("Updating Salesforce Object")
       updateSalesforceObject(username, password, login, version, sfObject.get, mode,
-          flag(upsert, "upsert"), externalIdFieldName, data)
+          flag(upsert, "upsert"), externalIdFieldName, batchSize, data)
     }
 
     return createReturnRelation(data)
@@ -156,6 +166,7 @@ class DefaultSource extends RelationProvider with SchemaRelationProvider with Cr
       mode: SaveMode,
       upsert: Boolean,
       externalIdFieldName: String,
+      batchSize: Integer,
       data: DataFrame) {
 
     val csvHeader = Utils.csvHeadder(data.schema)
@@ -164,7 +175,7 @@ class DefaultSource extends RelationProvider with SchemaRelationProvider with Cr
     val repartitionedRDD = Utils.repartition(data.rdd)
     logger.info("no of partitions after repartitioning is " + repartitionedRDD.partitions.length)
 
-    val writer = new SFObjectWriter(username, password, login, version, sfObject, mode, upsert, externalIdFieldName, csvHeader)
+    val writer = new SFObjectWriter(username, password, login, version, sfObject, mode, upsert, externalIdFieldName, csvHeader, batchSize)
     logger.info("Writing data")
     val successfulWrite = writer.writeData(repartitionedRDD)
     logger.info(s"Writing data was successful was $successfulWrite")

src/main/scala/com/springml/spark/salesforce/SFObjectWriter.scala

@@ -23,20 +23,24 @@ class SFObjectWriter (
     val mode: SaveMode,
     val upsert: Boolean,
     val externalIdFieldName: String,
-    val csvHeader: String
+    val csvHeader: String,
+    val batchSize: Integer
     ) extends Serializable {
 
   @transient val logger = Logger.getLogger(classOf[SFObjectWriter])
 
   def writeData(rdd: RDD[Row]): Boolean = {
     val csvRDD = rdd.map(row => row.toSeq.map(value => Utils.rowValue(value)).mkString(","))
 
+    val partitionCnt = (1 + csvRDD.count() / batchSize).toInt
+    val partitionedRDD = csvRDD.repartition(partitionCnt)
+
     val jobInfo = new JobInfo(WaveAPIConstants.STR_CSV, sfObject, operation(mode, upsert))
     jobInfo.setExternalIdFieldName(externalIdFieldName)
 
     val jobId = bulkAPI.createJob(jobInfo).getId
 
-    csvRDD.mapPartitionsWithIndex {
+    partitionedRDD.mapPartitionsWithIndex {
       case (index, iterator) => {
         val records = iterator.toArray.mkString("\n")
         var batchInfoId : String = null