Add parameter for partitionByKey

[ZekLouis]

Hey, we're using your package and we need to precise on which key our partitions will be based, that's why I made this change :)

Tell me if you need more explanations about our use case :)

[Rydgel]

I have the same problem, this MR is welcomed 👍

[Yolooman]

Same here 👍

[gaets86]

Really helpful for me too