Pyspark Linting Rules

[sbrugman]

Apache Spark is widely used in the python ecosystem for distributed computing. As user of spark I would like for ruff to lint problematic behaviours. The automation that ruff offers is especially useful in projects with various levels of software engineering skills, e.g. where people has more of a statistics background.

There exists a pyspark style guide and pylint extension.

I would like to start contributing a rule that checks for repeated use of withColumn:

This method introduces a projection internally. Therefore, calling it multiple times, for instance, via loops in order to add multiple columns can generate big plans which can cause performance issues and even StackOverflowException. To avoid this, use select() with multiple columns at once.

Are you ok with a PR introducing "Spark-specific rules" (e.g. SPK)?

ruff includes rules that are specific to third party libraries: numpy, pandas and airflow. Spark support would be a nice addition.

I would like to close with the following thought: supporting third-party packages may at first seem to be effort in the long tail of possible rules to add to ruff. Why not focus only on rules that affect all Python users? I hope that adding these will lead to creating helper functions that make adding new rules easier. I also think that these libraries will end up with similar API design patterns, that can be linted across the ecosystem. As an example, call chaining is common for many packages that perform transformations.

[charliermarsh]

I'm generally open to adding package-specific rule sets for extremely popular packages (as with Pandas, NumPy, etc.), and Spark would fit that description. However, it'd be nice to have a few rules lined up before we move forward and add any one of them. Otherwise, we run the risk that we end up with really sparse categories that only contain a rule or two.

[sbrugman]

Super. I've updated the issue with a couple of rules that we can track.

[guilhem-dvr]

Hi, I was looking for such a thread.

To add to the proposed list, here are some rules we wish we had at my company:

• unnecessary drop followed by a select
• use unionByName instead of union / unionAll
• use df.writeTo(...).append() instead of df.write.insertInto(...)
• use df.writeTo(...).overwritePartitions() instead of df.write.insertInto(..., overwrite=True)
• replace udf with native spark functions
• alias pyspark.sql.functions to F -> from pyspark.sql import ..., functions as F, ...

[amadeuspzs]

Just to add that I would be interested in this functionality.

Also, the first link in the original post is broken, and the pylint extension looks unmaintained?

• https://github.com/palantir/pyspark-style-guide is a written PySpark guide, and contains some pylint implementations under https://github.com/palantir/pyspark-style-guide/tree/develop/src/checkers
• https://nhsdigital.github.io/rap-community-of-practice/training_resources/pyspark/pyspark-style-guide/ is based off the guide above

[stkrzysiak]

Super. I've updated the issue with a couple of rules that we can track. I'll kick off with SPK001-3

Did you get going with this? Thinking about jumping on it

[Rexeh]

Currently looking at recommending Ruff for data team and this feature would be great.

[aran159]

Super. I've updated the issue with a couple of rules that we can track. I'll kick off with SPK001-3

Did you get going with this? Thinking about jumping on it

Hey! Did you end up jumping on it? We're also considering starting it, so I'd love to hear how it went for you!

[sbrugman]

I'm working on getting a first set of rules out there :)

[sbrugman]

@guilhem-dvr Thanks for your suggestions! Would you have an example for the following?

unnecessary drop followed by a select

Also note that the import convention is already possible via: https://docs.astral.sh/ruff/settings/#lint_flake8-import-conventions_aliases

[guilhem-dvr]

Sure, here's what I had in mind:

df = spark.createDataFrame(
    [("John", 25, "Engineer"), ("Jane", 30, "Doctor"), ("Jim", 35, "Teacher")],
    ["Name", "Age", "Profession"],
)

# Uncessary drop
df.drop("Age").select("Name", "Profession")

# Same statement without the drop
df.select("Name", "Profession")

But now I see that there's a pattern that shouldn't be flagged: where an 'anti' select is performed, i.e. drop some columns then select all the remaining ones with select("*"):

# Reusing the previous df schema
df.drop("Age").select("*")

Edit: this is still a bad pattern because drop already returns the whole dataframe - minus the dropped column - so you should never be chaining drop and select anyway.

Also note that the import convention is already possible via: https://docs.astral.sh/ruff/settings/#lint_flake8-import-conventions_aliases

Lol, I had completely skimmed over the settings, thank you!

[sbrugman]

Thanks for clarifying! There is not as much Spark open-source code available as there is for other libraries, so it's hard to tell how frequent this error is - but a good addition nonetheless. Funny enough, Polars uses this in their tests: https://github.com/pola-rs/polars/blob/main/py-polars/tests/unit/lazyframe/test_lazyframe.py#L1090

[stkrzysiak]

Super. I've updated the issue with a couple of rules that we can track. I'll kick off with SPK001-3

Did you get going with this? Thinking about jumping on it

Hey! Did you end up jumping on it? We're also considering starting it, so I'd love to hear how it went for you!

No, but looks like @sbrugman has, which is exciting!