chore(typing): enable typing checks for `pyspark` by EdAbati · Pull Request #2051 · narwhals-dev/narwhals

EdAbati · 2025-02-20T07:35:54Z

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

Related issues

Checklist

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

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

for more information, see https://pre-commit.ci

tests/frame/unpivot_test.py

for more information, see https://pre-commit.ci

dangotbanned · 2025-02-24T09:28:21Z

@EdAbati do we want this merged before or after #2038?

narwhals/_spark_like/expr.py

narwhals/_spark_like/dataframe.py

narwhals/_spark_like/namespace.py

dangotbanned · 2025-02-24T10:25:59Z

I'd really like this PR to be resolving warnings vs just adding ignore comments.
I do understand we don't have stubs to work with, so @EdAbati let me know if this is an unreasonable ask

EdAbati · 2025-02-24T12:04:21Z

Thanks for the comments, I'll move the unpivot fix out so we can concentrate on discussing typing here. I'll reply later

https://github.com/narwhals-dev/narwhals/actions/runs/13526995704/job/37799975130?pr=2051

dangotbanned · 2025-02-25T17:38:32Z

https://github.com/narwhals-dev/narwhals/actions/runs/13527240491/job/37800780977?pr=2051

This one is gonna take a little extra work to fix

Update

(11bde2e) is a pretty fast-and-loose solution to the problem.

https://github.com/EdAbati/narwhals/blob/11bde2e95c93d10f4cc04a389a259da2a05d1f6b/narwhals/_spark_like/dataframe.py#L43

We should think about moving the Implementation-dispatch into two subclasses:

Keep all the shared behavior in SparkLikeLazyFrame
Make that generic over TypeVar("T", DataFrame, SQLFrameDataFrame)
Move _implementation into a ClassVar on each subclass
Then those subclasses just have the unique parts defined
- native_dataframe, _F, _native_dtypes, _Window, etc

Important

Food for thought - not proposing these changes become part of this PR

Fixes this and all the issues that propogated from solving https://github.com/narwhals-dev/narwhals/actions/runs/13527240491/job/37800780977?pr=2051

dangotbanned

Thanks for working on this PR @EdAbati!

Had a lot of fun spiraling through the issues that came up here.
I'll leave this open to let someone else merge - since I've made a lot of changes that I'd also want reviewed 😄

FBruzzesi

I must admit I still get easily lost in the typing shenanigans - can't spot anything that does not sound reasonable, but it's not a review that you should trust either 🙈

FBruzzesi · 2025-02-25T19:17:14Z

narwhals/_spark_like/dataframe.py

    from narwhals.dtypes import DType
    from narwhals.utils import Version

+    SQLFrameDataFrame: TypeAlias = _SQLFrameDataFrame[Any, Any, Any, Any, Any]


Damn that's a hell of a generic 😂

Damn that's a hell of a generic 😂

Ikr!
All the Unknown(s) in this gave me quite the scare
https://github.com/narwhals-dev/narwhals/actions/runs/13527240491/job/37800780977?pr=2051

FBruzzesi · 2025-02-25T19:26:08Z

narwhals/_spark_like/dataframe.py

+        if TYPE_CHECKING:
+            from pyspark.sql import types
+
+            return types


Is there any drawback to this? I have barely seen these conditional in non-import blocks

Is there any drawback to this? I have barely seen these conditional in non-import blocks

@FBruzzesi

Oh let me introduce you to how pytest works then 😄

https://github.com/pytest-dev/pytest/blob/9cf2cae94355cd83ad7e8d88f976c5a524c98cfb/src/_pytest/mark/structures.py#L432-L485
https://github.com/pytest-dev/pytest/blob/9cf2cae94355cd83ad7e8d88f976c5a524c98cfb/src/_pytest/mark/structures.py#L505-L512
https://github.com/pytest-dev/pytest/blob/9cf2cae94355cd83ad7e8d88f976c5a524c98cfb/src/_pytest/_py/path.py#L203-L217

polars does some pretty complex stuff here - which includes tricking the type checker to recognise these symbols & getting lazy imports from them

https://github.com/pola-rs/polars/blob/45ec22a18c0cbdc67f56b3230f8b76149679ef02/py-polars/polars/dependencies.py
https://github.com/pola-rs/polars/blob/45ec22a18c0cbdc67f56b3230f8b76149679ef02/py-polars/polars/dataframe/frame.py#L88-L90

Is there any drawback to this? I have barely seen these conditional in non-import blocks

I'm not sure I'd say drawback - but I would recommend only trying this as a last resort.
Here - I'm using it to describe something the type system doesn't support yet.

Another use was #2064 (comment) - for a backwards compat bug fix.

I'd much rather do things in a simple way whenever possible

I did a similar thing here to convert these functions into typeguards

https://github.com/zen-xu/pyarrow-stubs/blob/c482d8f146a8cf2cff55cc4060148fe4883f9309/pyarrow-stubs/types.pyi

narwhals/narwhals/_arrow/utils.py

Lines 18 to 62 in 81353a3

if TYPE_CHECKING:

from typing import TypeVar

from typing_extensions import TypeAlias

from typing_extensions import TypeIs

from narwhals._arrow.series import ArrowSeries

from narwhals._arrow.typing import ArrowArray

from narwhals._arrow.typing import ArrowChunkedArray

from narwhals._arrow.typing import Incomplete

from narwhals._arrow.typing import StringArray

from narwhals.dtypes import DType

from narwhals.typing import _AnyDArray

from narwhals.utils import Version

# NOTE: stubs don't allow for `ChunkedArray[StructArray]`

# Intended to represent the `.chunks` property storing `list[pa.StructArray]`

ChunkedArrayStructArray: TypeAlias = ArrowChunkedArray

_T = TypeVar("_T")

def is_timestamp(t: Any) -> TypeIs[pa.TimestampType[Any, Any]]: ...

def is_duration(t: Any) -> TypeIs[pa.DurationType[Any]]: ...

def is_list(t: Any) -> TypeIs[pa.ListType[Any]]: ...

def is_large_list(t: Any) -> TypeIs[pa.LargeListType[Any]]: ...

def is_fixed_size_list(t: Any) -> TypeIs[pa.FixedSizeListType[Any, Any]]: ...

def is_dictionary(

t: Any,

) -> TypeIs[pa.DictionaryType[Any, Any, Any]]: ...

def extract_regex(

strings: ArrowChunkedArray,

/,

pattern: str,

*,

options: Any = None,

memory_pool: Any = None,

) -> ChunkedArrayStructArray: ...

else:

from pyarrow.compute import extract_regex

from pyarrow.types import is_dictionary # noqa: F401

from pyarrow.types import is_duration

from pyarrow.types import is_fixed_size_list

from pyarrow.types import is_large_list

from pyarrow.types import is_list

from pyarrow.types import is_timestamp

Forgot to mention, but (#2051 (comment)) would be the direction to go in - to avoid needing TYPE_CHECKING blocks.

Keeping everything in one class means you need to either do this - or introduce Union(s) and isinstance checks to resolve them on every usage - to type this correctly.

Currently, we branch on Implementation but we can't represent the relationship between Implementation and these unions:

pyspark.sql.types or sqlframe.base.types

pyspark.sql.functions or sqlframe.base.functions

pyspark.sql.Window or sqlframe.base.window.Window

pyspark.sql.dataframe.DataFrame.sparkSession or sqlframe.base.dataframe.BaseDataFrame.session

They're a mix of modules, classes and object attributes

MarcoGorelli

from a quick-ish scan this looks good, but i didn't check the details, though if you're all happy with it feel free to ship it 🚢

dangotbanned · 2025-02-27T22:14:56Z

from a quick-ish scan this looks good, but i didn't check the details, though if you're all happy with it feel free to ship it 🚢

Ah perfect thanks @MarcoGorelli yeah I'm happy with it if you are ♥️

Edit: ffs every time with the double gif 🤦

Porting over (#2051), didn't realize this was delclared twice until (#2132)

* chore(typing): Add typing for `SparkLikeExpr` properties Porting over (#2051), didn't realize this was delclared twice until (#2132) * chore: fix typing and simplify `SparkLikeExprStringNamespace.to_datetime` Resolves (https://github.com/narwhals-dev/narwhals/actions/runs/13682912007/job/38259412675?pr=2152) * rename function * use single chars in set * fix: Remove timezone offset replacement * test: Adds `test_to_datetime_tz_aware` Resolves #2152 (comment) * test: possibly fix `pyarrow` in ci? Maybe this was just a TZDATA issue locally? https://github.com/narwhals-dev/narwhals/actions/runs/13699734154/job/38310256617?pr=2152 * test: xfail polars `3.8`, fix false positive pyarrow https://github.com/narwhals-dev/narwhals/actions/runs/13699804987/job/38310487932?pr=2152 https://github.com/narwhals-dev/narwhals/actions/runs/13699804987/job/38310488783?pr=2152 * test: narrower xfail, tz-less expected? Not even sure what `pyarrow` is doing here https://github.com/narwhals-dev/narwhals/actions/runs/13700021595/job/38311197947?pr=2152 * test: account for `pyarrow` version changes https://github.com/narwhals-dev/narwhals/actions/runs/13700267075/job/38312036397?pr=215 * test: maybe fix `pyspark` https://github.com/narwhals-dev/narwhals/actions/runs/13700361438/job/38312364899?pr=2152 * revert: go back to typing fixes only Addresses #2152 (review) * chore: ignore `format` shadowing #2152 (review) * keep logic the same I hope #2152 (comment) --------- Co-authored-by: Edoardo Abati <29585319+EdAbati@users.noreply.github.com>

- Added in narwhals-dev#2187, narwhals-dev#2051 - No longer needed since `sqlframe` only typing

EdAbati and others added 8 commits February 16, 2025 13:23

remove pyspark from mypy overrides

9c420d0

some mypy fixes

cc2b5de

Merge remote-tracking branch 'upstream/main' into pyspark-typing

59a93f5

Merge remote-tracking branch 'upstream/main' into pyspark-typing

0a33350

ignore type var

f2c8e6a

add more tests

e429278

we already test on = none

8588444

[pre-commit.ci] auto fixes from pre-commit.com hooks

803cad9

for more information, see https://pre-commit.ci

EdAbati commented Feb 20, 2025

View reviewed changes

tests/frame/unpivot_test.py Outdated Show resolved Hide resolved

EdAbati and others added 8 commits February 24, 2025 08:46

merge upstream

2e8c9d0

[pre-commit.ci] auto fixes from pre-commit.com hooks

68107d1

for more information, see https://pre-commit.ci

fix pyspark unpivot

9558c69

Merge branch 'main' into pyspark-typing

a3740d4

add pyspark typing ci

53c10ee

fix type

f4f3ec7

simplify type unpivot

c71e3f9

make pyright happy

5b4ce3a

EdAbati mentioned this pull request Feb 24, 2025

[Bug]: mismatch behaviour with .unpivot(on=None, index=None) #2082

Open

Merge branch 'main' into pyspark-typing

37861aa

EdAbati marked this pull request as ready for review February 24, 2025 09:31

dangotbanned reviewed Feb 24, 2025

View reviewed changes

narwhals/_spark_like/expr.py Outdated Show resolved Hide resolved

narwhals/_spark_like/expr.py Outdated Show resolved Hide resolved

narwhals/_spark_like/dataframe.py Outdated Show resolved Hide resolved

narwhals/_spark_like/namespace.py Outdated Show resolved Hide resolved

dangotbanned added the typing label Feb 24, 2025

dangotbanned mentioned this pull request Feb 24, 2025

refactor: Generic CompliantSelector #2064

Merged

19 tasks

EdAbati added 3 commits February 25, 2025 08:38

Merge remote-tracking branch 'upstream/main' into pyspark-typing

6ff24c2

fix invert type

e6c08c8

ignore return any

a353ff3

EdAbati changed the title ~~fix: pyspark typing and related unpivot bug~~ fix: enable typing check for pyspark Feb 25, 2025

fix(typing): resolve [attr-defined]

9448db0

https://github.com/narwhals-dev/narwhals/actions/runs/13526995704/job/37799975130?pr=2051

dangotbanned changed the title ~~fix: enable typing checks for pyspark~~ chore(typing): enable typing checks for pyspark Feb 25, 2025

Merge remote-tracking branch 'upstream/main' into pr/EdAbati/2051

d04af00

dangotbanned marked this pull request as draft February 25, 2025 17:38

fix(typing): temp resolve impl differences

11bde2e

Fixes this and all the issues that propogated from solving https://github.com/narwhals-dev/narwhals/actions/runs/13527240491/job/37800780977?pr=2051

dangotbanned marked this pull request as ready for review February 25, 2025 18:10

dangotbanned approved these changes Feb 25, 2025

View reviewed changes

dangotbanned requested review from FBruzzesi and MarcoGorelli February 25, 2025 18:14

dangotbanned added the internal label Feb 25, 2025

FBruzzesi reviewed Feb 25, 2025

View reviewed changes

dangotbanned added 3 commits February 26, 2025 10:06

Merge branch 'main' into pyspark-typing

93ccbfc

Merge branch 'main' into pyspark-typing

10ebdc1

Merge branch 'main' into pyspark-typing

271f726

MarcoGorelli approved these changes Feb 27, 2025

View reviewed changes

dangotbanned merged commit 5383756 into narwhals-dev:main Feb 27, 2025
27 of 28 checks passed

dangotbanned added a commit that referenced this pull request Mar 5, 2025

chore(typing): Add typing for SparkLikeExpr properties

47e768d

Porting over (#2051), didn't realize this was delclared twice until (#2132)

dangotbanned mentioned this pull request Mar 5, 2025

chore(typing): SparkLikeExpr properties #2152

Merged

10 tasks

dangotbanned mentioned this pull request Mar 11, 2025

typing: use SQLFrame instead of PySpark for typing #2185

Closed

dangotbanned added a commit to MarcoGorelli/narwhals that referenced this pull request Mar 12, 2025

chore(typing): remove unused ignores

4a1aabf

- Added in narwhals-dev#2187, narwhals-dev#2051 - No longer needed since `sqlframe` only typing

This was referenced Mar 31, 2025

chore: Add Compliant*.from_native #2315

Merged

Make import_optional_dependency util #1940

Open

This was referenced Apr 18, 2025

feat: let nw.Enum accept categories, map pandas ordered categorical to Enum (only in main namespace, not stable.v1) #2192

Merged

feat: Add support for spark connect #2417

Merged

This was referenced Aug 27, 2025

refactor: Organizing Compliant* APIs #3045

Merged

Transition to single-Implementation classes #3059

Open

dangotbanned mentioned this pull request Jan 21, 2026

support a Module type in ty_extensions for typing specific modules astral-sh/ty#2574

Closed

	if TYPE_CHECKING:
	from typing import TypeVar

	from typing_extensions import TypeAlias
	from typing_extensions import TypeIs

	from narwhals._arrow.series import ArrowSeries
	from narwhals._arrow.typing import ArrowArray
	from narwhals._arrow.typing import ArrowChunkedArray
	from narwhals._arrow.typing import Incomplete
	from narwhals._arrow.typing import StringArray
	from narwhals.dtypes import DType
	from narwhals.typing import _AnyDArray
	from narwhals.utils import Version

	# NOTE: stubs don't allow for `ChunkedArray[StructArray]`
	# Intended to represent the `.chunks` property storing `list[pa.StructArray]`
	ChunkedArrayStructArray: TypeAlias = ArrowChunkedArray

	_T = TypeVar("_T")

	def is_timestamp(t: Any) -> TypeIs[pa.TimestampType[Any, Any]]: ...
	def is_duration(t: Any) -> TypeIs[pa.DurationType[Any]]: ...
	def is_list(t: Any) -> TypeIs[pa.ListType[Any]]: ...
	def is_large_list(t: Any) -> TypeIs[pa.LargeListType[Any]]: ...
	def is_fixed_size_list(t: Any) -> TypeIs[pa.FixedSizeListType[Any, Any]]: ...
	def is_dictionary(
	t: Any,
	) -> TypeIs[pa.DictionaryType[Any, Any, Any]]: ...
	def extract_regex(
	strings: ArrowChunkedArray,
	/,
	pattern: str,
	*,
	options: Any = None,
	memory_pool: Any = None,
	) -> ChunkedArrayStructArray: ...
	else:
	from pyarrow.compute import extract_regex
	from pyarrow.types import is_dictionary # noqa: F401
	from pyarrow.types import is_duration
	from pyarrow.types import is_fixed_size_list
	from pyarrow.types import is_large_list
	from pyarrow.types import is_list
	from pyarrow.types import is_timestamp

Conversation

EdAbati commented Feb 20, 2025 • edited by dangotbanned Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

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

Related issues

Checklist

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

Uh oh!

Uh oh!

dangotbanned commented Feb 24, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

dangotbanned commented Feb 24, 2025

Uh oh!

EdAbati commented Feb 24, 2025

Uh oh!

dangotbanned commented Feb 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Update

Uh oh!

dangotbanned left a comment

Choose a reason for hiding this comment

Uh oh!

FBruzzesi left a comment

Choose a reason for hiding this comment

Uh oh!

FBruzzesi Feb 25, 2025

Choose a reason for hiding this comment

Uh oh!

dangotbanned Feb 25, 2025

Choose a reason for hiding this comment

Uh oh!

FBruzzesi Feb 25, 2025

Choose a reason for hiding this comment

Uh oh!

dangotbanned Feb 25, 2025

Choose a reason for hiding this comment

Uh oh!

dangotbanned Feb 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

FBruzzesi Feb 25, 2025

Choose a reason for hiding this comment

Uh oh!

dangotbanned Feb 25, 2025

Choose a reason for hiding this comment

Uh oh!

dangotbanned Feb 25, 2025

Choose a reason for hiding this comment

Uh oh!

dangotbanned Feb 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

MarcoGorelli left a comment

Choose a reason for hiding this comment

Uh oh!

dangotbanned commented Feb 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

EdAbati commented Feb 20, 2025 •

edited by dangotbanned

Loading

dangotbanned commented Feb 25, 2025 •

edited

Loading

dangotbanned Feb 25, 2025 •

edited

Loading

dangotbanned Feb 26, 2025 •

edited

Loading

dangotbanned commented Feb 27, 2025 •

edited

Loading