Fixes the binary math functions:

• atan2 and hypot take two arguments, not one
• pow supports taking a literal numeric value as its second argument in addition to a Column.

It seems like there is some problem with Jedi compatibility. Some components seem to work pretty well. For example DataFrame without stubs:

In [1]: import jedi                                                                                                                                                                                                

In [2]: from pyspark.sql import SparkSession                                                                                                                                                                       

In [3]: jedi.Interpreter("SparkSession.builder.getOrCreate().createDataFrame([]).", [globals()]).completions()                                                                                                     
---------------------------------------------------------------------------
AttributeError   
...
AttributeError: 'ModuleContext' object has no attribute 'py__path__'

and with stubs:

In [1]: from pyspark.sql import SparkSession                                                                                                                                                                       

In [2]: import jedi                                                                                                                                                                                                

In [3]: jedi.Interpreter("SparkSession.builder.getOrCreate().createDataFrame([]).", [globals()]).completions()                                                                                                     
Out[3]: 
[<Completion: agg>,
 <Completion: alias>,
 <Completion: approxQuantile>,
 <Completion: cache>,
 <Completion: checkpoint>,
 <Completion: coalesce>,
 <Completion: collect>,
 <Completion: colRegex>,
 <Completion: columns>,
 <Completion: corr>,
 <Completion: count>,
 <Completion: cov>,
...
 <Completion: __str__>]

So far so good. However, if take for example LinearRegressionModel.load things don't work so well. Without stubs provides no suggestions

In [1]: import jedi                                                                                                                                                                                                

In [2]: from pyspark.ml.regression import LinearRegressionModel                                                                                                                                                    

In [3]: jedi.Interpreter("LinearRegressionModel.load('foo').", [globals()]).completions()                                                                                                                          
Out[3]: []

but one provided with stubs

In [1]: import jedi                                                                                                                                                                                                

In [2]: from pyspark.ml.regression import LinearRegressionModel                                                                                                                                                    

In [3]: jedi.Interpreter("LinearRegressionModel.load('foo').", [globals()]).completions()                                                                                                                          
Out[3]: 
[<Completion: load>,
 <Completion: read>,
 <Completion: __annotations__>,
 <Completion: __class__>,
 <Completion: __delattr__>,
 <Completion: __dict__>,
 <Completion: __dir__>,
 <Completion: __doc__>,
 <Completion: __eq__>,
 <Completion: __format__>,
 <Completion: __getattribute__>,
 <Completion: __hash__>,
 <Completion: __init__>,
 <Completion: __init_subclass__>,
 <Completion: __module__>,
 <Completion: __ne__>,
 <Completion: __new__>,
 <Completion: __reduce__>,
 <Completion: __reduce_ex__>,
 <Completion: __repr__>,
 <Completion: __setattr__>,
 <Completion: __sizeof__>,
 <Completion: __slots__>,

don't make much sense. If model is fitted:

In [4]: from pyspark.ml.regression import LinearRegression                                                                                                                                                         

In [5]: jedi.Interpreter("LinearRegression().fit(...).", [globals()]).completions()                                                                                                                                
Out[5]: 
[<Completion: aggregationDepth>,
 <Completion: append>,
 <Completion: clear>,
 <Completion: coefficients>,
 <Completion: copy>,
 <Completion: count>,
....
 <Completion: __str__>]

Model which is explicitly annotated works fine, so it seems like there is something in MLReader or one of the sub-classes that causes a failure.

We already have data tests for this (as well as some test cases from apache/spark examples, and mypy seems to be fine with this.

Since LinearRegression.fit works fine (and some toy tests confirm that), Generics are not sufficient to reproduce the problem. So it seems like type parameter is not processed correctly on the path:

• LinearRegressionModel(..., avaMLReadable[LinearRegressionModel], ...)
• JavaMLReadable(MLReadable[R])
• MLReadable(Generic[R])
• load(cls, path: str) -> R

Tested with:

• jedi==0.15.2 and jedi==0.16.0 (0c56aa4).
• pyspark-stubs==3.0.0.dev5
• pyspark==3.0.0.dev0 (afe70b3)

Using 2.4.0.post6

spark.read.load(folders, inferSchema=True, header=False)

mypy reports Expected type 'str', got 'bool' instead for both inferSchema and header.

Looks like the issue is in third_party/3/pyspark/sql/readwriter.pyi Line 23 where in the definition for load() we have **options: str. For csv suppport this needs to be **options: Optional[Union[bool, str, int]] but to handle the general case it probably needs to be **options: Any.

The contains method is missing from the stubs causing mypy to raise error: "Column" not callable.

This PR adds the typehints to 2.4 specifically (the version we are using), but it should probably also be added to the other versions.

Passing a List[str] to select raises a mypy warning, similar for List[Column]. We change the type from List[Union[Column, str]] to Union[List[Column], List[str]].

Fixes #394 .

Update repartition functions to allow for Col in numPartitions parameter.

Reference

numPartitions – can be an int to specify the target number of partitions or a Column.
    If it is a Column, it will be used as the first partitioning column.
    If not specified, the default number of partitions is used.

Also add stub for DataFrame#distinct()

In the functions here: https://github.com/zero323/pyspark-stubs/blob/3c4684a224c1be4eea4577e475f8bb4d045edddd/third_party/3/pyspark/sql/functions.pyi#L100-L101 we currently have tz: str but this can also be specified as a Column

Example:

>>> from pyspark.sql import functions
>>> df = spark.sql("SELECT CAST(0 AS TIMESTAMP) AS timestamp, 'Asia/Tokyo' AS tz")
>>> df.select(functions.from_utc_timestamp(df.timestamp, df.tz)).collect()
[Row(from_utc_timestamp(timestamp, tz)=datetime.datetime(1970, 1, 1, 18, 0))]

I think this could be expanded to tz: ColumnOrName?

The method DataFrame.drop expects either 1 Column, or 1 str, or an iterable of strings. This is only type checked inside the function though.

Currently the type hints (and the actual API) allow to pass multiple Columns but it does result in a runtime error. Personally, I'd like to have that caught earlier. But as this might be getting too close to the internals of the functions, I’d like to hear your opinion on whether or not the type hints should “look inside” to aid development.

The fraction is a required argument to the sample method. Anytime someone calls df.sample(.01) this is met in mypy with

Argument 1 to "sample" of "DataFrame" has incompatible type "float"; expected "Optional[bool]"

In the Pyspark API, the three arguments are in fact pure keyword arguments that are handled later to ensure fraction must be given. This is probably done to keep consistent with the Scala API.

By overloading the methods, the issue is resolved.

Resolves #273

Additional parameters to DataFrameReader.load() and DataFrameWriter.save()/.saveTable() are passed to the file-type specific reader or writer types. These parameters can be of any type.

2.3 has these data types and I was erroneously gettting errors for them.

Note this is a port of e2d225f06ff36fcbf79e2123f1c18f380e862728

I tried a cherry-pick but it had some issues (not sure why)