[SPARK-3418] New code review comments addressed