typing: use SQLFrame instead of PySpark for typing

[MarcoGorelli]

SQLFrame have added a py.typed file in their latest release

Could we use them for typing instead of PySpark? The API should be the same so it should provide the same benefits code-navigation-wise, but wouldn't require contributors to install pyspark

[dangotbanned]

IIRC we'd need both because of the different import paths.

I haven't checked the diff from (https://github.com/eakmanrq/sqlframe/releases/tag/v3.24.0) yet - but the previous version only had similar (not exact) annotations which might be a hurdle if we tried it out

  /home/runner/work/narwhals/narwhals/narwhals/_spark_like/dataframe.py:286:45 - error: Argument of type "Column" cannot be assigned to parameter "column" of type "Column | str | bool"
    Type "Column" is not assignable to type "Column | str | bool"
      "pyspark.sql.column.Column" is not assignable to "sqlframe.base.column.Column"
      "Column" is not assignable to "str"
      "Column" is not assignable to "bool" (reportArgumentType)

Related

• chore(typing): enable typing checks for pyspark #2051 (comment)
• chore(typing): enable typing checks for pyspark #2051 (comment)
• https://github.com/narwhals-dev/narwhals/actions/runs/13785020372/job/38550899134#step:9:32

[MarcoGorelli]

maybe for _NativeDataFrame, but for the rest (e.g. Window, functions, Column) could we just use the sqlframe one?

[dangotbanned]

maybe for _NativeDataFrame, but for the rest (e.g. Window, functions, Column) could we just use the sqlframe one?

I totally get the motivation for this @MarcoGorelli - I too would like a no-pyspark option!

The problem is they're all connected.

For example, in #2187 I'm trying out using a constrained TypeVar (must be A or B, but exclusive).
But that isn't enough on it's own - because in some places we do use just one type.
Here it is pyspark.sql.column.Column- but the same would be true for the inverse (as you've suggested)

---

I can most likely get this to work with the current SparkLike* classes.
But IMO should really just be a stepping stone to #2051 (comment)
The solution will still require casting in certain places that we branch on implementation and do anything non-trivial like:

narwhals/narwhals/_spark_like/dataframe.py

Lines 439 to 478 in 01830c1

	if self._implementation.is_pyspark():
	return self._from_native_frame(
	native_frame.select(
	*[
	self._F.col(col_name).alias(col_name)
	if col_name != columns[0]
	else self._F.explode_outer(col_name).alias(col_name)
	for col_name in column_names
	]
	),
	)
	elif self._implementation.is_sqlframe():
	# Not every sqlframe dialect supports `explode_outer` function
	# (see https://github.com/eakmanrq/sqlframe/blob/3cb899c515b101ff4c197d84b34fae490d0ed257/sqlframe/base/functions.py#L2288-L2289)
	# therefore we simply explode the array column which will ignore nulls and
	# zero sized arrays, and append these specific condition with nulls (to
	# match polars behavior).
	def null_condition(col_name: str) -> Column:
	return self._F.isnull(col_name) | (self._F.array_size(col_name) == 0)
	return self._from_native_frame(
	native_frame.select(
	*[
	self._F.col(col_name).alias(col_name)
	if col_name != columns[0]
	else self._F.explode(col_name).alias(col_name)
	for col_name in column_names
	]
	).union(
	native_frame.filter(null_condition(columns[0])).select(
	*[
	self._F.col(col_name).alias(col_name)
	if col_name != columns[0]
	else self._F.lit(None).alias(col_name)
	for col_name in column_names
	]
	)
	),
	)

[MarcoGorelli]

sure, but we don't have many such cases, right?

like in _pandas_like we type everything as if it were pandas, even though it also handles Modin and cuDF. I was thinking we could do the same there in _spark_like

[dangotbanned]

sure, but we don't have many such cases, right?

I've got 12 at the moment I think - but that might go up or down while fixing 😅

like in _pandas_like we type everything as if it were pandas, even though it also handles Modin and cuDF. I was thinking we could do the same there in _spark_like

Ah this reminded me of (#2051 (comment)) which never got resolved

@FBruzzesi

@MarcoGorelli I forgot what was the reason for not returning SQLFrame modules from base - wouldn't SQLFrame itself take care of dispatching to the engine? That would avoid having to add a dedicated block for type checking

There seems to be some functionality in SQLFrame that might simplify stuff for us https://sqlframe.readthedocs.io/en/stable/configuration/#activating-sqlframe

pinging @eakmanrq in case that might helps us

[MarcoGorelli]

It should be fine to use things in sqlframe.base for annotations, but we can't set anything global (e.g. sqlframe.activate, or pd.options copy-on-write) on behalf of the user

[dangotbanned]

It should be fine to use things in sqlframe.base for annotations, but we can't set anything global (e.g. sqlframe.activate, or pd.options copy-on-write) on behalf of the user

This should be fine though right - if its scoped in a with block?
https://sqlframe.readthedocs.io/en/stable/configuration/#context-based-activation

Note

Still doesn't resolve the install issue though

[MarcoGorelli]

sorry what's the install issue?

if we use sqlframe for functions / Window / Column, then we should still be able to type check the internal _spark_like code without pyspark installed, and if pyspark is used in a TYPE_CHECKING block and we don't have it installed, then wouldn't it just get ignored by the type checker?

[dangotbanned]

sorry what's the install issue?

Ah apologies I wasn't clear

@MarcoGorelli

Could we use them for typing instead of PySpark? The API should be the same so it should provide the same benefits code-navigation-wise, but wouldn't require contributors to install pyspark

I meant that the lower install cost benefit you wanted - wouldn't be something we could get from using my suggestion of pyspark types but replacing things at runtime via sqlframe.
IIUC there needs to be something to replace first.

Although this has made me think about another issue with changing pyspark -> sqlframe as you suggested.
Wouldn't that break typing for existing users that use pyspark without sqlframe?

Edit

I think the inverse is currently true if you can install sqlframe without pyspark - that those users will run into typing issues

[dangotbanned]

if we use sqlframe for functions / Window / Column, then we should still be able to type check the internal _spark_like code without pyspark installed, and if pyspark is used in a TYPE_CHECKING block and we don't have it installed, then wouldn't it just get ignored by the type checker?

I guess it depends how they have their type checker configured - but I think if we go exclusively pyspark (current) or sqlframe (proposed) - at least one group of users could start seeing Unknown(s) in their code without both.

If that is a problem - it would be resolved by splitting the two into subclasses (per #2051 (comment)).
Then we can keep the promise of whatever you pass in you get back

[MarcoGorelli]

this isn't user-facing, are we sure it would matter to users?