feat(datasets): Add support for databricks-connect>=13.0

[MigQ2]

Description

Closes #351

Checklist

• Updated the documentation to reflect the code changes
• Added a description of this change in the relevant RELEASE.md file
• Added tests to cover my changes (I don't think there's an available Databricks environment to test on)

[cdkrot]

For even further improved compatibility, add

from pyspark.sql.connect.session import SparkSession as RemoteSparkSession

and replace this line with RemoteSparkSession.builder.getOrCreate().

[MigQ2]

Thanks a lot for looking into this @cdkrot , super appreciated

Do you mean another try except? So the fallback waterfall would be:

1. DatabricksSession
2. RemoteSparkSession
3. SparkSession

I guess adding RemoteSparksession would helo adding support for non-Databricks Spark-connect sessions?

[cdkrot]

Sorry, my bad, the logic in code is good as is 👍. sql.SparkSession will figure out if it needs to create sql.SparkSession or connect.SparkSession.

Possibly it's best to rewrite the comments though. Essentially the only reason to fallback from dbconnect session creation to pyspark's is when first is not installed, since dbconnect switches to pyspark automatically based on environment.

[MigQ2]

@astrojuanlu could you please re-run the CI/CD checks to see if the current implementation passes all checks and is ready to merge?

[astrojuanlu]

Thanks for the ping @MigQ2 ! Hope everything is good to go now. Feel free to ping me again if needed.

[astrojuanlu]

And also thanks @cdkrot for reviewing, much appreciated 🙌🏽

[MigQ2]

Hi @astrojuanlu , @noklam I think I need some help with figuring out the best way to keep spark tests working fine here.

I see there's some wrangling with mocker.patch and _get_spark so the way I built it breaks the tests as each dataset wouldn't have its own _get_spark() method.

I have thought about the following workaround but I'd like to know what you think before implementing it:

We could move the logic of _get_spark() to a utils.py (or keep it in spark_dataset.py, if you prefer that), and keep a self._get_spark() staticmethod in all DataSet classes which would be just a wrapper for utils._get_spark().

This way we could keep tests as they are because each DataSet class would have its own self._get_spark() method and I think mocker.patch would work fine as it is withhout changing the test

WDYT?

[noklam]

I see there's some wrangling with mocker.patch and _get_spark so the way I built it breaks the tests as each dataset wouldn't have its own _get_spark() method.

In this case we could update the test to mock get_spark instead of mocker.object

We could move the logic of _get_spark() to a utils.py (or keep it in spark_dataset.py, if you prefer that), and keep a self._get_spark() staticmethod in all DataSet classes which would be just a wrapper for utils._get_spark().

I think it's fine to keep it in spark_dataset.py, if we ever move it I may prefer it stay in the spark/__init__.py instead of another new utils.py.

[merelcht]

Hi @MigQ2 , are you able to address the review comments soon? We'd like to do a kedro-datasets release this week and it would be great to include your contribution if possible. Otherwise, it will be part of our next release.

[MigQ2]

Hi @merelcht , I'd love to see it on the next release so I'll work on it ASAP.

Is there an easy way I can run the test suite and CI/CD locally to verify everything is working fine without pinging someone from the kedro team to manually trigger it im GitHub?

[astrojuanlu]

@MigQ2 Can you try make test plugin=kedro-datasets or even cd kedro-datasets && pytest -k databricks?

[MigQ2]

Thanks @astrojuanlu , I managed to run tests locally and they should succeed now. Can you pls trigger CI/CD again to see if everything is alright?

I struggled a bit to get stuff working, so here go some ideas that maybe could help some wannabe-kedroids to contribute more easily:

1. Document how to run tests (basically document the cd kedro-datasets && pytest -k databricks you just mentioned before), I found it by myself for the kedro core, but I didn't find it so obvious for kedro-plugins

2. Help to setup the local environment to run the tests. I had to install all dependencies using pip install ./kedro-datasets/[test], which took a while because of the big requirements list (especially pyspark and tensorflow).
   And then I realized I had to install java8 (which is a bit hard because it's been removed from package manager) for spark tests to run locally.
   I like using devcontainers in vscode to have a consistent environment. Do you think it would be a good idea to provide a devcontainer with all dependencies (pypi, java, other system libraries, etc.) installed out of the box so that people can just open their repo in that container and run tests easily without worrying about installing all stuff in their system? I have a Dockerfile that worked for me, if you think it's a good idea we can talk and I can share

[noklam]

@MigQ2 We have Gitpod which is similar to codes pace and you can have it on browser, local Pycharm or VScode.

[MigQ2]

There seemed to be a dask conflict in library requirements. I guess that's not coming from my changes.

LMK when that's fixed and I'll update the branch and let's re-trigger the CI/CD

[merelcht]

There seemed to be a dask conflict in library requirements. I guess that's not coming from my changes.

LMK when that's fixed and I'll update the branch and let's re-trigger the CI/CD

That issue has been resolved! Now the test builds are failing because of missing coverage:
kedro_datasets/spark/spark_dataset.py 186 1 99% 46

[MigQ2]

Thanks! I'll check it out tonight and push some changes to fix coverage

[DimedS]

Thanks for the PR, @MigQ2. Can you clarify a few points for me? If I'm understanding this correctly, this PR enables remote connections to Databricks, prioritising the Databricks connector in all scenarios. However, when it comes to remote Spark session connections, is there a different syntax? For instance, would this work for a remote connection to an AWS EMR Spark cluster?

[MigQ2]

Thanks for the PR, @MigQ2. Can you clarify a few points for me? If I'm understanding this correctly, this PR enables remote connections to Databricks, prioritising the Databricks connector in all scenarios. However, when it comes to remote Spark session connections, is there a different syntax? For instance, would this work for a remote connection to an AWS EMR Spark cluster?

Hi @DimedS,

I haven't used standard spark-connnect sessions, and I don't have an AWS EMR or similar setup to test on.
As per spark-connect documentation I think syntax would be different, and something like the following code should be used:

from pyspark.sql import SparkSession
spark = SparkSession.builder.remote("some-host").getOrCreate()

This PR only attempts to fix for Databricks setups, cause that's the one we use and the one we can test on, but I think it would be great to add support for non-databricks remote spark-connect sessions

[MigQ2]

@merelcht I added # pragma: no cover to the databricks-connect session as I don't think there's a Databricks environment available to test on.

When the release is out I'll start using it with Databricks in our environment and come back here if I find anything that needs to be fixed os that we can improve

[MigQ2]

As a side comment, I tried out gitpod to setup the environment but there were a couple things that still weren't clear for me:

1. When running make install-test-requirements plugin=kedro-datasets I still got a dask dependency error
2. I still don't know what's the recommended way to install java8 in the gitpod environment so that pyspark test can run properly

[DimedS]

Thanks for the PR, @MigQ2. Can you clarify a few points for me? If I'm understanding this correctly, this PR enables remote connections to Databricks, prioritising the Databricks connector in all scenarios. However, when it comes to remote Spark session connections, is there a different syntax? For instance, would this work for a remote connection to an AWS EMR Spark cluster?

Hi @DimedS,

I haven't used standard spark-connnect sessions, and I don't have an AWS EMR or similar setup to test on. As per spark-connect documentation I think syntax would be different, and something like the following code should be used:

from pyspark.sql import SparkSession
spark = SparkSession.builder.remote("some-host").getOrCreate()

This PR only attempts to fix for Databricks setups, cause that's the one we use and the one we can test on, but I think it would be great to add support for non-databricks remote spark-connect sessions

yes, looks like this, that we need to add additional issue for making possible also remote spark connects.
It looks unclear that our default choice will be trying to connect to Databricks in any case, but I think for now it's ok to fastly allow Databricks remote connection. Maybe later we should rewrite it with more general solution.

[MigQ2]

Thanks for the PR, @MigQ2. Can you clarify a few points for me? If I'm understanding this correctly, this PR enables remote connections to Databricks, prioritising the Databricks connector in all scenarios. However, when it comes to remote Spark session connections, is there a different syntax? For instance, would this work for a remote connection to an AWS EMR Spark cluster?

Hi @DimedS,
I haven't used standard spark-connnect sessions, and I don't have an AWS EMR or similar setup to test on. As per spark-connect documentation I think syntax would be different, and something like the following code should be used:

from pyspark.sql import SparkSession
spark = SparkSession.builder.remote("some-host").getOrCreate()

This PR only attempts to fix for Databricks setups, cause that's the one we use and the one we can test on, but I think it would be great to add support for non-databricks remote spark-connect sessions

yes, looks like this, that we need to add additional issue for making possible also remote spark connects. It looks unclear that our default choice will be trying to connect to Databricks in any case, but I think for now it's ok to fastly allow Databricks remote connection. Maybe later we should rewrite it with more general solution.

Yes, agree. I think the try-except is not that bad cause I don't think anyone would install the databricks-connect package in a non-databricks setup.

Another point that I think would be nice is to think about remote sessions authentication. For now this PR doesn't have any authentication options (or any other spark options), so it can only authenticate with environment variables.

However there are several other different ways to authenticate to databricks (and probably to spark-connect) that maybe other users would like to use and this implementation wouldn't support

[noklam]

@MigQ2 can you open a separate issue for GitPod? I may have time to look at that, pretty sure this was working before we migrate to kedro-plugins so I can find it in older commits.

[DimedS]

@MigQ2, thank you for your contribution. It enables the use of remote Spark sessions with Databricks. However, we should keep the initial issue open, as this PR doesn't facilitate remote Spark sessions with solutions other than Databricks.