scala - Spark : DB connection per Spark RDD partition and do mapPartition

Question

I want to do a mapPartitions on my spark rdd,

    val newRd = myRdd.mapPartitions(
      partition => {

        val connection = new DbConnection /*creates a db connection per partition*/

        val newPartition = partition.map(
           record => {
             readMatchingFromDB(record, connection)
         })
        connection.close()
        newPartition
      })

But, this gives me a connection already closed exception, as expected because before the control reaches the .map() my connection is closed. I want to create a connection per RDD partition, and close it properly. How can I achieve this?

Thanks!

Answer 1

As mentioned in the discussion here - the issue stems from the laziness of map operation on the iterator partition. This laziness means that for each partition, a connection is created and closed, and only later (when RDD is acted upon), readMatchingFromDB is called.

To resolve this, you should force an eager traversal of the iterator before closing the connection, e.g. by converting it into a list (and then back):

val newRd = myRdd.mapPartitions(partition => {
  val connection = new DbConnection /*creates a db connection per partition*/

  val newPartition = partition.map(record => {
    readMatchingFromDB(record, connection)
  }).toList // consumes the iterator, thus calls readMatchingFromDB 

  connection.close()
  newPartition.iterator // create a new iterator
})