egen · mosche · Sep 15, 2020
diff --git a/README.md b/README.md
@@ -43,26 +43,30 @@ $ bin/spark-shell --packages com.springml:spark-salesforce_2.11:1.1.3
 * `username`: Salesforce Wave Username. This user should have privilege to upload datasets or execute SAQL or execute SOQL  
 * `password`: Salesforce Wave Password. Please append security token along with password.For example, if a user’s password is mypassword, and the security token is XXXXXXXXXX, the user must provide mypasswordXXXXXXXXXX
 * `login`: (Optional) Salesforce Login URL. Default value https://login.salesforce.com
-* `datasetName`: (Optional) Name of the dataset to be created in Salesforce Wave. Required for Dataset Creation
-* `sfObject`: (Optional) Salesforce Object to be updated. (e.g.) Contact. Mandatory if `bulk` is `true`.
+* `sfObject`: (Optional) Salesforce Object to be fetched or updated. (e.g.) Contact. Mandatory if `bulk` is `true`.
 * `metadataConfig`: (Optional) Metadata configuration which will be used to construct [Salesforce Wave Dataset Metadata] (https://resources.docs.salesforce.com/sfdc/pdf/bi_dev_guide_ext_data_format.pdf). Metadata configuration has to be provided in JSON format
 * `saql`: (Optional) SAQL query to used to query Salesforce Wave. Mandatory for reading Salesforce Wave dataset
 * `soql`: (Optional) SOQL query to used to query Salesforce Object. Mandatory for reading Salesforce Object like Opportunity
 * `version`: (Optional) Salesforce API Version. Default 35.0
-* `inferSchema`: (Optional) Inferschema from the query results. Sample rows will be taken to find the datatype
+* `inferSchema`: (Optional) Infer schema from the query results. Sample rows will be taken to find the datatype
 * `dateFormat`: (Optional) A string that indicates the format that follow java.text.SimpleDateFormat to use when reading timestamps. This applies to TimestampType. By default, it is null which means trying to parse timestamp by java.sql.Timestamp.valueOf()
 * `resultVariable`: (Optional) result variable used in SAQL query. To paginate SAQL queries this package will add the required offset and limit. For example, in this SAQL query `q = load \"<dataset_id>/<dataset_version_id>\"; q = foreach q generate  'Name' as 'Name',  'Email' as 'Email';` **q** is the result variable
 * `pageSize`: (Optional) Page size for each query to be executed against Salesforce Wave. Default value is 2000. This option can only be used if `resultVariable` is set
-* `upsert`: (Optional) Flag to upsert data to Salesforce. This performs an insert or update operation using the "externalIdFieldName" as the primary ID. Existing fields that are not in the dataframe being pushed will not be updated. Default "false".
 
 ### Options only supported for fetching Salesforce Objects.
 * `bulk`: (Optional) Flag to enable bulk query. This is the preferred method when loading large sets of data. Salesforce will process batches in the background. Default value is `false`.
 * `pkChunking`: (Optional) Flag to enable automatic primary key chunking for bulk query job. This splits bulk queries into separate batches that of the size defined by `chunkSize` option. By default `false` and the default chunk size is 100,000.
 * `chunkSize`: (Optional) The size of the number of records to include in each batch. Default value is 100,000. This option can only be used when `pkChunking` is `true`. Maximum size is 250,000.
 * `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 writing Salesforce Objects / Salesforce Wave datasets.
+* `datasetName`: (Optional) Name of the dataset to be created in Salesforce Wave. Required for Dataset Creation
+* `upsert`: (Optional) Flag to upsert data to Salesforce. This performs an insert or update operation using the `externalIdFieldName` as the primary ID. Existing fields that are not in the dataframe being pushed will not be updated. Default "false".
+* `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".
+* `batchSize`: (Optional) The size in bytes per ingest batch. Default value is `10 MB`, maximum size is `10 MB`.
+* `batchRecords`: (Optional) The number of records per ingest batch. Per default only `batchSize` is considered. Maximum number of records per batch is `10,000`.
+
 
 ### Scala API
 ```scala

diff --git a/src/main/scala/com/springml/spark/salesforce/DefaultSource.scala b/src/main/scala/com/springml/spark/salesforce/DefaultSource.scala
@@ -18,6 +18,7 @@ package com.springml.spark.salesforce
 import java.text.SimpleDateFormat
 
 import com.springml.salesforce.wave.api.APIFactory
+import org.apache.commons.io.FileUtils.byteCountToDisplaySize
 import org.apache.http.Header
 import org.apache.http.message.BasicHeader
 import org.apache.log4j.Logger
@@ -26,6 +27,7 @@ import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.{DataFrame, SQLContext, SaveMode}
 
 import scala.collection.mutable.ListBuffer
+import scala.util.Try
 
 /**
  * Default source for Salesforce wave data source.
@@ -61,29 +63,25 @@ class DefaultSource extends RelationProvider with SchemaRelationProvider with Cr
     val saql = parameters.get("saql")
     val soql = parameters.get("soql")
     val resultVariable = parameters.get("resultVariable")
-    val pageSize = parameters.getOrElse("pageSize", "1000")
-    val sampleSize = parameters.getOrElse("sampleSize", "1000")
-    val maxRetry = parameters.getOrElse("maxRetry", "5")
-    val inferSchema = parameters.getOrElse("inferSchema", "false")
+    val pageSize = getIntParam(parameters, "pageSize").getOrElse(1000)
+    val sampleSize = getIntParam(parameters, "sampleSize").getOrElse(1000)
+    val maxRetry = getIntParam(parameters, "maxRetry").getOrElse(5)
+    val inferSchemaFlag = getBooleanParam(parameters, "inferSchema").getOrElse(false)
     val dateFormat = parameters.getOrElse("dateFormat", null)
     // This is only needed for Spark version 1.5.2 or lower
     // Special characters in older version of spark is not handled properly
     val encodeFields = parameters.get("encodeFields")
     val replaceDatasetNameWithId = parameters.getOrElse("replaceDatasetNameWithId", "false")
 
-    val bulkStr = parameters.getOrElse("bulk", "false")
-    val bulkFlag = flag(bulkStr, "bulk")
-
-    val queryAllStr = parameters.getOrElse("queryAll", "false")
-    val queryAllFlag = flag(queryAllStr, "queryAll")
+    val bulkFlag = getBooleanParam(parameters, "bulk").getOrElse(false)
+    val queryAllFlag = getBooleanParam(parameters, "queryAll").getOrElse(false)
 
     validateMutualExclusive(saql, soql, "saql", "soql")
-    val inferSchemaFlag = flag(inferSchema, "inferSchema")
 
     if (saql.isDefined) {
       val waveAPI = APIFactory.getInstance.waveAPI(username, password, login, version)
       DatasetRelation(waveAPI, null, saql.get, schema, sqlContext,
-          resultVariable, pageSize.toInt, sampleSize.toInt,
+          resultVariable, pageSize, sampleSize,
           encodeFields, inferSchemaFlag, replaceDatasetNameWithId.toBoolean, sdf(dateFormat),
           queryAllFlag)
     } else {
@@ -98,8 +96,7 @@ class DefaultSource extends RelationProvider with SchemaRelationProvider with Cr
       if (bulkFlag) {
         createBulkRelation(sqlContext, username, password, login, version, inferSchemaFlag, parameters, schema)
       } else {
-        val forceAPI = APIFactory.getInstance.forceAPI(username, password, login,
-          version, Integer.getInteger(pageSize), Integer.getInteger(maxRetry))
+        val forceAPI = APIFactory.getInstance.forceAPI(username, password, login, version, pageSize, maxRetry)
         DatasetRelation(null, forceAPI, soql.get, schema, sqlContext,
           null, 0, sampleSize.toInt, encodeFields, inferSchemaFlag,
           replaceDatasetNameWithId.toBoolean, sdf(dateFormat), queryAllFlag)
@@ -118,16 +115,16 @@ class DefaultSource extends RelationProvider with SchemaRelationProvider with Cr
     val login = parameters.getOrElse("login", "https://login.salesforce.com")
     val version = parameters.getOrElse("version", "36.0")
     val usersMetadataConfig = parameters.get("metadataConfig")
-    val upsert = parameters.getOrElse("upsert", "false")
+    val upsertFlag = getBooleanParam(parameters, "upsert").getOrElse(false)
+    val batchSize = getLongParam(parameters, "batchSize").getOrElse(1024 * 1024 * 10L)
+    val batchRecords = getIntParam(parameters, "batchRecords")
     val metadataFile = parameters.get("metadataFile")
-    val encodeFields = parameters.get("encodeFields")
-    val monitorJob = parameters.getOrElse("monitorJob", "false")
+    val monitorJobFlag = getBooleanParam(parameters, "monitorJob").getOrElse(false)
     val externalIdFieldName = parameters.getOrElse("externalIdFieldName", "Id")
 
     validateMutualExclusive(datasetName, sfObject, "datasetName", "sfObject")
 
     if (datasetName.isDefined) {
-      val upsertFlag = flag(upsert, "upsert")
       if (upsertFlag) {
         if (metadataFile == null || !metadataFile.isDefined) {
           sys.error("metadataFile has to be provided for upsert" )
@@ -137,14 +134,14 @@ class DefaultSource extends RelationProvider with SchemaRelationProvider with Cr
       logger.info("Writing dataframe into Salesforce Wave")
       writeInSalesforceWave(username, password, login, version,
           datasetName.get, appName, usersMetadataConfig, mode,
-          flag(upsert, "upsert"), flag(monitorJob, "monitorJob"), data, metadataFile)
+          upsertFlag, monitorJobFlag, batchSize, batchRecords, data, metadataFile)
     } else {
       logger.info("Updating Salesforce Object")
       updateSalesforceObject(username, password, login, version, sfObject.get, mode,
-          flag(upsert, "upsert"), externalIdFieldName, data)
+          upsertFlag, externalIdFieldName, batchSize, batchRecords, data)
     }
 
-    return createReturnRelation(data)
+    createReturnRelation(data)
   }
 
   private def updateSalesforceObject(
@@ -156,13 +153,15 @@ class DefaultSource extends RelationProvider with SchemaRelationProvider with Cr
       mode: SaveMode,
       upsert: Boolean,
       externalIdFieldName: String,
+      batchSize: Long,
+      batchRecords: Option[Int],
       data: DataFrame) {
 
     val csvHeader = Utils.csvHeadder(data.schema)
-    logger.info("no of partitions before repartitioning is " + data.rdd.partitions.length)
-    logger.info("Repartitioning rdd for 10mb partitions")
-    val repartitionedRDD = Utils.repartition(data.rdd)
-    logger.info("no of partitions after repartitioning is " + repartitionedRDD.partitions.length)
+    logger.info("Number of partitions before repartitioning is " + data.rdd.partitions.length)
+    logger.info(s"Repartitioning rdd for ${byteCountToDisplaySize(batchSize)} partitions${batchRecords.fold("")(n => s" with $n records")}")
+    val repartitionedRDD = Utils.repartition(data.rdd, batchSize, batchRecords)
+    logger.info("Number of partitions after repartitioning is " + repartitionedRDD.partitions.length)
 
     val writer = new SFObjectWriter(username, password, login, version, sfObject, mode, upsert, externalIdFieldName, csvHeader)
     logger.info("Writing data")
@@ -190,31 +189,14 @@ class DefaultSource extends RelationProvider with SchemaRelationProvider with Cr
       throw new Exception("sfObject must not be empty when performing bulk query")
     }
 
-    val timeoutStr = parameters.getOrElse("timeout", "600000")
-    val timeout = try {
-      timeoutStr.toLong
-    } catch {
-      case e: Exception => throw new Exception("timeout must be a valid integer")
-    }
+    val timeout = getLongParam(parameters, "timeout").getOrElse(600000L)
 
     var customHeaders = ListBuffer[Header]()
-    val pkChunkingStr = parameters.getOrElse("pkChunking", "false")
-    val pkChunking = flag(pkChunkingStr, "pkChunkingStr")
+    val pkChunking = getBooleanParam(parameters, "pkChunking").getOrElse(false)
 
     if (pkChunking) {
-      val chunkSize = parameters.get("chunkSize")
-
-      if (!chunkSize.isEmpty) {
-        try {
-          chunkSize.get.toInt
-        }
-        catch {
-          case e: Exception => throw new Exception("chunkSize must be a valid integer")
-        }
-        customHeaders += new BasicHeader("Sforce-Enable-PKChunking", s"chunkSize=${chunkSize.get}")
-      } else {
-        customHeaders += new BasicHeader("Sforce-Enable-PKChunking", "true")
-      }
+      val pkChunkingValue = getIntParam(parameters, "chunkSize").fold("true")(size => s"chunkSize=$size")
+      customHeaders += new BasicHeader("Sforce-Enable-PKChunking", pkChunkingValue)
     }
 
     BulkRelation(
@@ -243,6 +225,8 @@ class DefaultSource extends RelationProvider with SchemaRelationProvider with Cr
       mode: SaveMode,
       upsert: Boolean,
       monitorJob: Boolean,
+      batchSize: Long,
+      batchRecords: Option[Int],
       data: DataFrame,
       metadata: Option[String]) {
     val dataWriter = new DataWriter(username, password, login, version, datasetName, appName)
@@ -257,10 +241,10 @@ class DefaultSource extends RelationProvider with SchemaRelationProvider with Cr
     }
     logger.info(s"Able to write the metadata is $writtenId")
 
-    logger.info("no of partitions before repartitioning is " + data.rdd.partitions.length)
-    logger.info("Repartitioning rdd for 10mb partitions")
-    val repartitionedRDD = Utils.repartition(data.rdd)
-    logger.debug("no of partitions after repartitioning is " + repartitionedRDD.partitions.length)
+    logger.info("Number of partitions before repartitioning is " + data.rdd.partitions.length)
+    logger.info(s"Repartitioning rdd for ${byteCountToDisplaySize(batchSize)} partitions${batchRecords.fold("")(n => s" with $n records")}")
+    val repartitionedRDD = Utils.repartition(data.rdd, batchSize, batchRecords)
+    logger.debug("Number of partitions after repartitioning is " + repartitionedRDD.partitions.length)
 
     logger.info("Writing data")
     val successfulWrite = dataWriter.writeData(repartitionedRDD, writtenId.get)
@@ -322,13 +306,15 @@ class DefaultSource extends RelationProvider with SchemaRelationProvider with Cr
     simpleDateFormat
   }
 
-  private def flag(paramValue: String, paramName: String) : Boolean = {
-    if (paramValue == "false") {
-      false
-    } else if (paramValue == "true") {
-      true
-    } else {
-      sys.error(s"""'$paramName' flag can only be true or false""")
-    }
+  private def getBooleanParam(params: Map[String, String], paramName: String) = params.get(paramName).map {
+    case "false" => false
+    case "true" => true
+    case _ => sys.error(s"'$paramName' flag can only be true or false")
   }
+
+  private def getLongParam(params: Map[String, String], paramName: String) = params.get(paramName)
+    .map(v => Try(v.toLong).getOrElse(sys.error(s"'$paramName' must be of type Long")))
+
+  private def getIntParam(params: Map[String, String], paramName: String) = params.get(paramName)
+    .map(v => Try(v.toInt).getOrElse(sys.error(s"'$paramName' must be of type Int")))
 }
diff --git a/src/main/scala/com/springml/spark/salesforce/Utils.scala b/src/main/scala/com/springml/spark/salesforce/Utils.scala
@@ -24,12 +24,9 @@ import com.madhukaraphatak.sizeof.SizeEstimator
 import org.apache.log4j.Logger
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.Row
-import org.apache.spark.sql.types.{DoubleType, IntegerType, StructType}
+import org.apache.spark.sql.types.{StructType}
 
-import scala.collection.immutable.HashMap
 import com.springml.spark.salesforce.metadata.MetadataConstructor
-import com.sforce.soap.partner.sobject.SObject
-import scala.concurrent.duration._
 import com.sforce.soap.partner.fault.UnexpectedErrorFault
 
 import scala.concurrent.duration.FiniteDuration
@@ -62,34 +59,30 @@ object Utils extends Serializable {
     })
   }
 
-  def repartition(rdd: RDD[Row]): RDD[Row] = {
-    val totalDataSize = getTotalSize(rdd)
-    val maxBundleSize = 1024 * 1024 * 10l
-    var partitions = 1
-    if (totalDataSize > maxBundleSize) {
-      partitions = Math.round(totalDataSize / maxBundleSize) + 1
-    }
+  def repartition(rdd: RDD[Row], maxBundleSize: Long, maxBundleRecords: Option[Int]): RDD[Row] = {
+    val totalRows = rdd.count()
+    val totalDataSize = getTotalSize(rdd, totalRows)
+
+    val partitionsByBytes = if (totalDataSize > maxBundleSize) Math.round(totalDataSize / maxBundleSize) + 1 else 1
+    val partitionsByRecords = maxBundleRecords.fold(1)(rows => if (totalRows > rows) Math.round(totalRows / rows) + 1 else 1)
+
+    val partitions = Math.max(partitionsByBytes, partitionsByRecords)
 
     val shuffle = rdd.partitions.length < partitions
-    rdd.coalesce(partitions.toInt, shuffle)
+    rdd.coalesce(partitions, shuffle)
   }
 
-  def getTotalSize(rdd: RDD[Row]): Long = {
+  def getTotalSize(rdd: RDD[Row], totalRows: Long): Long = {
     // This can be fetched as optional parameter
     val NO_OF_SAMPLE_ROWS = 10
-    val totalRows = rdd.count()
-    var totalSize = 0l
 
     if (totalRows > NO_OF_SAMPLE_ROWS) {
       val sampleObj = rdd.takeSample(false, NO_OF_SAMPLE_ROWS)
       val sampleRowSize = rowSize(sampleObj)
-      totalSize = sampleRowSize * (totalRows / NO_OF_SAMPLE_ROWS)
+      sampleRowSize * (totalRows / NO_OF_SAMPLE_ROWS)
     } else {
-
-      totalSize = rddSize(rdd)
+      rddSize(rdd)
     }
-
-    totalSize
   }
 
   def rddSize(rdd: RDD[Row]) : Long = {

diff --git a/src/test/scala/com/springml/spark/salesforce/TestUtils.scala b/src/test/scala/com/springml/spark/salesforce/TestUtils.scala
@@ -81,25 +81,34 @@ class TestUtils extends FunSuite with BeforeAndAfterEach {
     val schema = StructType(columnStruct)
     val inMemoryDF = ss.sqlContext.createDataFrame(inMemoryRDD, schema)
 
-    val repartitionDF = Utils.repartition(inMemoryRDD)
+    val repartitionDF = Utils.repartition(inMemoryRDD, maxBundleSize = 1024 * 1024 * 10L, maxBundleRecords = None)
     assert(repartitionDF.partitions.length == 1)
   }
 
-  test("Test repartition for local CSV file with size less than 10 MB") {
+  test("Test repartition for local CSV file with size less than maxBundleSize") {
     val csvURL= getClass.getResource("/ad-server-data-formatted.csv")
     val csvFilePath = csvURL.getPath
     val csvDF = ss.read.option("header", "true").csv(csvFilePath)
 
-    val repartitionDF = Utils.repartition(csvDF.rdd)
+    val repartitionDF = Utils.repartition(csvDF.rdd, maxBundleSize = 1024 * 1024 * 10L, maxBundleRecords = None)
     assert(repartitionDF.partitions.length == 1)
   }
 
-  test("Test repartition for local CSV file with size > 10 MB and < 20 MB") {
+  test("Test repartition for local CSV file with size less than maxBundleSize but exceeding maxBundleRecords") {
+    val csvURL= getClass.getResource("/ad-server-data-formatted.csv")
+    val csvFilePath = csvURL.getPath
+    val csvDF = ss.read.option("header", "true").csv(csvFilePath)
+
+    val repartitionDF = Utils.repartition(csvDF.rdd, maxBundleSize = 1024 * 1024 * 10L, maxBundleRecords = Some(100))
+    assert(repartitionDF.partitions.length == 4)
+  }
+
+  test("Test repartition for local CSV file with size exceeding maxBundleSize") {
     val csvURL= getClass.getResource("/minified_GDS_90.csv")
     val csvFilePath = csvURL.getPath
     val csvDF = ss.read.option("header", "true").csv(csvFilePath)
 
-    val repartitionDF = Utils.repartition(csvDF.rdd)
+    val repartitionDF = Utils.repartition(csvDF.rdd, maxBundleSize = 1024 * 1024 * 10L, maxBundleRecords = None)
     assert(repartitionDF.partitions.length == 2)
   }