Backport SPARK-1124 fix

https://spark-project.atlassian.net/browse/SPARK-1124