[FEATURE REQUEST]: Support for Spark 3.1.1

Spark 3.1.1 release is coming soon. I ran E2E tests against the RC version of Spark 3.1.1 and the following two tests fail due to the internal changes, which we should address to support Spark 3.1.1:
```
TestDataFrameGroupedMapUdf
TestGroupedMapUdf
```
