Spark Structure Streaming - add a batch column with value from currentBatchId

Question

Currently, we are trying to add a new column to the incoming streaming dataframe with value from currentBatchId of StreamExecution. This is required by our use case.

Is that possible?

Answer 1

You can use the following UDF:

val batchUDF = udf { () =>
  TaskContext.get().getLocalProperty("streaming.sql.batchId").toInt
}

Then: df.withColumn("batch", batchUDF())