Make Ben's requested changes from PR https://github.com/kaskada-ai/kaskada/pull/799
Make Ben's requested changes from PR #799