[SPARK-3974] Code review feedback addressed