Spark (Scala) conditionally write JSON to multiple dynamic output locations - scala

Question

I have a problem to solve where data comes in as JSON from kinesis like below:
{
datatype: "datatype_1"
id : "id_1"
data : {...}
}
All records in the stream then need to go through a lookup function with datatype and id passed as arguments to find a unique group of locations to write the items to in JSON.
i.e.
def get_locations(id: String, datatype: String): Array[String] = //custom logic here
where the resultant array would look like
[ "s3:///example_bucket/example_folder_1", "s3:///example_bucket2/example_folder_2"]
My question is how do I most efficiently group records coming off the stream by datatype and id and write to the various s3 locations. I was hoping to do something like below:
sparkSession.readStream.format("kinesis")
.option("streamName", kinesis_stream_name)
.option("initialPosition", "latest")
.option("region", aws_region)
.load()
//more transforms
.select(
col("datatype"),
col("id"),
col("data")
)
// Not sure how I can do what's below
// .write.partitionBy("id", "datatype")
// .format("json")
// .option("compression","gzip")
// .save(get_locations("id","datatype")) //saving to all locations in result array

Answer 1

I do advise you to create the bucket in the code in the runtime as a best practice, you can use the node.js aws S3 API or your runtime language API
As you said in your comment you are getting the parameters from your runtime.
However as an answer to your question here is a function that creates a bucket containing the id in its name (you can change it to the format that you like ) then in that bucket you will have a lot of files based on the partition of the dataframe while saving:
import java.util
import com.amazonaws.regions.Regions
import com.amazonaws.services.s3.model.AmazonS3Exception
import com.amazonaws.services.s3.{AmazonS3, AmazonS3ClientBuilder}
def get_locations(id: String, datatype: String) = {
//you can configure the default region to the adequat region
//of course
val s3: AmazonS3 = AmazonS3ClientBuilder.standard.withRegion(Regions.DEFAULT_REGION).build
object CreateBucket {
def getBucket(bucket_name: String): Bucket = {
var named_bucket = null.asInstanceOf[Bucket]
val buckets: util.List[Bucket] = s3.listBuckets
import scala.collection.JavaConversions._
for (b <- buckets) {
if (b.getName.equals(bucket_name)) named_bucket = b
}
named_bucket
}
def createBucket(bucket_name: String): Bucket = {
var b = null.asInstanceOf[Bucket]
if (s3.doesBucketExistV2(bucket_name)) {
System.out.format("Bucket %s already exists.\n", bucket_name)
b = getBucket(bucket_name)
}
else try b = s3.createBucket(bucket_name)
catch {
case e: AmazonS3Exception =>
System.err.println(e.getErrorMessage)
}
b
}
}
//change your bucket name here if
//you like
val bucket_name = "bucket_" + id
var my_bucket = null.asInstanceOf[Bucket]
if (s3.doesBucketExistV2(bucket_name)) {
System.out.format("Bucket %s already exists.\n", bucket_name)
my_bucket = CreateBucket.getBucket(bucket_name)
}
else try my_bucket = s3.createBucket(bucket_name)
catch {
case e: AmazonS3Exception =>
System.err.println(e.getErrorMessage)
}
my_bucket
}
//I don't know how you will get those parameters
var id = " "
var datatype = " "
df.write.partitionBy("id", "dataType")
.format("json")
.option("compression", "gzip")
.save(get_locations(id, datatype).toString)
Don't forget to add the dependecies in maven or in build.sbt with the version that you have already in aws (sdk) :
<dependency>
<groupId>com.amazonaws</groupId>
<artifactId>aws-java-sdk-s3</artifactId>
<version>1.11.979</version>
</dependency>