Split spark dataset based on number of rows

I'm reading data from dynamo db and storing it in Spark Dataset like this :

// Building a dataset
            Dataset citations = sparkSession.read()
                .option("tableName", "Covid19Citation")
                .option("region", "eu-west-1")
                .format("dynamodb")
                .load();

What I want is to split this dataset based on number of rows.

For example, if dataset has more than 500 rows, I want to split it and save each of the datasets which I get as separate csv file. So, each of dataset which I want to save should have maximum 500 rows. Eg. If I there was 1600 rows in database, the output shoud be four xml files :

First xml file contains 500 rows,

Second xml file which contains also 500 rows,

Third xml file which contains also 500 rows, and finally

Fourth xml file which contains 100 rows.

This is what I tried so far, but this doesn't work:

List<Dataset> datasets = new ArrayList<>();
            while (citations.count() > 0) {
                  Dataset splitted = citations.limit(400);
                  datasets.add(splitted);
                  citations = citations.except(splitted);
            }

            System.out.println("datasets : " + datasets.size());
            for (Dataset d : datasets) {
                  code
                  d.coalesce(1)
                      .write()
                      .format("com.databricks.spark.xml")
                      .option("rootTag", "citations")
                      .option("rowTag", "citation")
                      .mode("overwrite")
                      .save("s3a://someoutputfolder/");
            }

Any help would be highly appreciated.

Thanks

You can leverage :

• row_number and mod: to split the dataset into parts of 500
• repartition: to produce one file per partition
• partitionBy: to write one xml per partition

Here an example in scala / parquet (but you can use xml as well)

val citations = spark.range(1, 2000000).selectExpr("id", "hash(id) value")

// calculate the number of buckets
val total = citations.count
val mod = (total.toFloat / 500).ceil.toInt

citations
.withColumn("id", expr("row_number() over(order by monotonically_increasing_id())"))
.withColumn("bucket", expr(f"mod(id, ${mod})"))
.repartition('bucket)
.write
.partitionBy("bucket")
.format("parquet")
.mode("overwrite")
.save("/tmp/foobar")


// now check the results
val resultDf = spark.read.format("parquet").load("/tmp/foobar")

// as a result you get at most 500 rows
scala> resultDf.groupBy("bucket").count.show
+------+-----+
|bucket|count|
+------+-----+
|  1133|  500|
|  1771|  500|
|  1890|  500|
|  3207|  500|
|  3912|  500|
|  1564|  500|
|  2823|  500|
+------+-----+

// there is no file with more than 500 rows
scala> resultDf.groupBy("bucket").count.filter("count > 500").show
+------+-----+
|bucket|count|
+------+-----+
+------+-----+

// now check there is only one file per bucket
scala> spark.sparkContext.parallelize(resultDf.inputFiles).toDF
.withColumn("part", expr("regexp_extract(value,'(bucket=([0-9]+))')"))
.groupBy("part").count.withColumnRenamed("count", "nb_files")
.orderBy(desc("nb_files")).show(5)
+-----------+--------+
|       part|nb_files|
+-----------+--------+
|bucket=3209|       1|
|bucket=1290|       1|
|bucket=3354|       1|
|bucket=2007|       1|
|bucket=2816|       1|
+-----------+--------+