FWIW this is what I ended up writing. I side-stepped the issue I was having by delegating the spark dataset read to PySpark: https://gist.github.com/williamcaicedo/f5379a668fe0f59f5dcd02f57bffa369