如何在spark数据帧v1.6中生成运行序列号

Question

我使用spark v1.6。我有以下数据框。

Primary_key | Dim_id

PK1 | 1

PK2 | 2

PK3 | 3

每当有新记录进入时，我想创建一个带有新序列#s的新数据帧。让我们说，我从源代码获得2条新记录，值为PK4＆amp; PK5，我想用值4和5创建新的dim_ids。所以，我的新数据框应该如下所示。

Primary_key | Dim_id

PK1 | 1

PK2 | 2

PK3 | 3

PK4 | 4

PK5 | 5

如何在spark数据帧v1.6中为新记录生成运行序列号？

Answer 1

如果你有一个数据库，你可以在其中创建一个序列，并将其与用户定义的函数一起使用（因为你，我偶然发现了这个问题......）。

保留一桶序列号并使用它（incrementby参数必须与用于创建序列的参数相同）。由于它是一个对象，SequenceID将是每个工作节点上的单例，您可以使用atomiclong迭代序列桶。

它远非完美（可能的连接泄漏，依赖于数据库，锁定静态类，确实如此），欢迎评论。

import java.sql.Connection
import java.sql.DriverManager
import java.util.concurrent.locks.ReentrantLock
import java.util.concurrent.atomic.AtomicLong
import org.apache.spark.sql.functions.udf

object SequenceID {
  var current: AtomicLong = new AtomicLong
  var max: Long = 0
  var connection: Connection = null
  var connectionLock = new ReentrantLock
  var seqLock = new ReentrantLock
  def getConnection(): Connection = {
    if (connection != null) {
      return connection
    }
    connectionLock.lock()
    if (connection == null) {
     // create your jdbc connection here
    }
    connectionLock.unlock()
    connection
  }

  def next(sequence: String, incrementBy: Long): Long = {
    if (current.get == max) {
      // sequence bucket exhausted, get a new one
      seqLock.lock()
      if (current.get == max) {
        val rs = getConnection().createStatement().executeQuery(s"SELECT NEXT VALUE FOR ${sequence} FROM sysibm.sysdummy1")
        rs.next()
        current.set(rs.getLong(1))
        max = current.get + incrementBy
      }
      seqLock.unlock()
    }
    return current.getAndIncrement
  }
}

class SequenceID() extends Serializable {
  def next(sequence: String, incrementBy: Long): Long = {
    return SequenceID.next(sequence, incrementBy)
  }  
}


val sequenceGenerator = new SequenceID(properties)

def sequenceUDF(seq: SequenceID) = udf[Long](() => {
  seq.next("PK_SEQUENCE", 500L)
})

val seq = sequenceUDF(sequenceGenerator)

myDataframe.select(myDataframe("foo"), seq())