Test repartitioning behaviour when writing parquet data.

test_repartition.scala

import java.util.UUID.randomUUID
import scala.sys.process._
import com.mozilla.telemetry.utils.getOrCreateSparkSession
import java.util.zip.CRC32

val spark = getOrCreateSparkSession("test")
spark.sparkContext.setLogLevel("WARN")
import spark.implicits._

def getSampleId(clientId: String, modulus: Int) = {
  // TODO: handle null
  val crc = new CRC32
  crc.update(clientId.getBytes)
  crc.getValue % modulus
}

val sampleIdModulus = 100
val filesPerPartition = 4

val sampleId = (c: String) => getSampleId(c, sampleIdModulus)
spark.sqlContext.udf.register("sampleid", sampleId)

val partitionId = (c: String) => getSampleId(c, filesPerPartition * sampleIdModulus)
spark.sqlContext.udf.register("partid", partitionId)

val df = (1 to 200000).map { e => ((e%10000).toString, randomUUID().toString) }.toList.toDF("client_id", "document_id")
val df2 = df.selectExpr("*", "sampleid(client_id) as sample_id", "partid(client_id) as part_id")

df2.repartition(df2.col("part_id")).drop("part_id").write.partitionBy("sample_id").mode("overwrite").parquet("/tmp/test-output")

"find /tmp/test-output -type f -name *.parquet" #| "wc -l" !

// 396 files. Four of the sample_id prefixes have only three files. CRC collisions maybe?