feat: Add minimal PySpark support

[EdAbati]

What type of PR is this? (check all applicable)

• 💾 Refactor
• ✨ Feature
• 🐛 Bug Fix
• 🔧 Optimization
• 📝 Documentation
• ✅ Test
• 🐳 Other

Related issues

• Closes [Enh]: Add Support For PySpark #333

Checklist

• Code follows style guide (ruff)
• Tests added
• Documented the changes

If you have comments or can explain your changes, please do so below.

As mentioned in the latest call, I've started working on the support for PySpark.

The goal of this PR would be to have a minimal initial implementation as a starting point. As we did for Dask, we can implement single methods in following PRs!

⚠️ this is not ready for review, a lot of test are failing, the code is ugly. :) Just making the PR for visibility and to have a place to comment/ask questions on specific points

[EdAbati]

This PR diff is getting big because of all the xfail in tests. 😕
@MarcoGorelli @FBruzzesi do you have a better idea on how to make it more "reviewable"? or do you think it is fine?

[FBruzzesi]

This PR diff is getting big because of all the xfail in tests. 😕 @MarcoGorelli @FBruzzesi do you have a better idea on how to make it more "reviewable"? or do you think it is fine?

For dask we started with it's own test file, so that we didn't have to modify every other file.
Once we had a few methods implemented we shifted the constructor into the conftest list of constructors and added the xfails.

Would that be a good strategy again?

[MarcoGorelli]

Seriously awesome and exciting work here, I'm impressed, well done!

Just left some comments / questions

Also, just noting something which I noticed:

(Pdb) p df.with_columns(d=nw.col('a').mean()).collect().to_native()
*** pyspark.errors.exceptions.captured.AnalysisException: [MISSING_GROUP_BY] The query does not include a GROUP BY clause. Add GROUP BY or turn it into the window functions using OVER clauses.;
Aggregate [a#1L, b#2L, c#3L, avg(a#1L) AS d#62]
+- Project [a#1L, b#2L, c#3L]
   +- Sort [index#0L ASC NULLS FIRST], true
      +- Repartition 2, true
         +- LogicalRDD [index#0L, a#1L, b#2L, c#3L], false

but we can deal with it later - I presume it shouldn't be too bad?

Can we also add a check in

narwhals/tests/no_imports_test.py

Lines 31 to 45 in 80aad57

	def test_pandas(monkeypatch: pytest.MonkeyPatch) -> None:
	monkeypatch.delitem(sys.modules, "polars")
	monkeypatch.delitem(sys.modules, "pyarrow")
	monkeypatch.delitem(sys.modules, "dask", raising=False)
	monkeypatch.delitem(sys.modules, "ibis", raising=False)
	df = pd.DataFrame({"a": [1, 1, 2], "b": [4, 5, 6]})
	nw.from_native(df, eager_only=True).group_by("a").agg(nw.col("b").mean()).filter(
	nw.col("a") > 1
	)
	assert "polars" not in sys.modules
	assert "pandas" in sys.modules
	assert "numpy" in sys.modules
	assert "pyarrow" not in sys.modules
	assert "dask" not in sys.modules
	assert "ibis" not in sys.modules

that in test_pandas or test_polars that 'pyspark' isn't in sys.modules at the end of the test?

[MarcoGorelli]

out of interest, is the str here just for typing purposes?

[EdAbati]

yes :)

[MarcoGorelli]

is there any time_unit / time_zone we should pass to dtypes.Datetime?

[EdAbati]

great point!

I think we cannot get it from the type itself but I need some more time to do some research.

If found things pyspark does when converting to pandas that could be useful to us:

• https://github.com/apache/spark/blob/master/python/pyspark/sql/pandas/types.py#L716
• https://github.com/apache/spark/blob/master/python/pyspark/sql/pandas/types.py#L607

Still have to spend some time to try to understand, but with a brief look it seems the time_unit is set to ns and time_zone is set to the local timezone.

For arrow it seems to be different though https://github.com/apache/spark/blob/master/python/pyspark/sql/pandas/types.py#L66

I don't have an opinion at the moment. What do you think?

[MarcoGorelli]

sure, happy to defer this, not a blocker

[MarcoGorelli]

amazing, thanks for updating! happy to merge once it's green as 🥦

[EdAbati]

Thank you all for the feedback and apologies this is taking longer but I am a bit busy.

The CI is almost green 😅 I need to add some if-else because pyspark is not available in python>3.11. Hopefully it will be fixed tonight

regarding the issue with df.with_columns(d=nw.col('a').mean()).collect().to_native()
I may have an idea on how to do it. I try to squeeze this is in too (and add a test). If this takes longer, could it be a follow-up?

[MarcoGorelli]

If this takes longer, could it be a follow-up?

yup definitely!

[EdAbati]

Ok CI is 🌲🌳

Regarding that issue, I am still working on a fix 🥲

[MarcoGorelli]

thanks @EdAbati for this massive effort

cool, let's ship it before merge conflicts creep it, excited about building this up!

[EdAbati]

Exciting 🎉🎉🎉 thanks!