[WIP] [SPARK-1698]: Improve spark integration