improve the code of SPARK-2096 [SQL]