Fix 2-argument math functions

third_party/3/pyspark/sql/_typing.pyi

@@ -1,4 +1,4 @@
-from typing import Any, List, Optional, TypeVar, Union
+from typing import Any, List, Optional, TypeVar, Union, SupportsFloat
 from typing_extensions import Protocol
 import datetime
 import decimal
@@ -7,6 +7,7 @@ import pyspark.sql.column
 import pyspark.sql.types
 
 ColumnOrName = Union[pyspark.sql.column.Column, str]
+ColumnNameOrFloat = Union[ColumnOrName, SupportsFloat]
 DecimalLiteral = decimal.Decimal
 DateTimeLiteral = Union[datetime.datetime, datetime.date]
 LiteralType = Union[bool, int, float, str]

third_party/3/pyspark/sql/functions.pyi

@@ -1,7 +1,7 @@
 # Stubs for pyspark.sql.functions (Python 3.5)
 
 from typing import overload
-from typing import Any, Optional, Union, Dict, Callable
+from typing import Any, Optional, Union, Dict, Callable, SupportsFloat
 
 import pandas.core.frame   # type: ignore
 import pandas.core.series     # type: ignore
@@ -164,7 +164,12 @@ def asc_nulls_last(col: ColumnOrName) -> Column: ...
 def ascii(col: ColumnOrName) -> Column: ...
 def asin(col: ColumnOrName) -> Column: ...
 def atan(col: ColumnOrName) -> Column: ...
-def atan2(col: ColumnOrName) -> Column: ...
+@overload
+def atan2(col1: ColumnOrName, col2: ColumnOrName) -> Column: ...
+@overload
+def atan2(col1: SupportsFloat, col2: ColumnOrName) -> Column: ...
+@overload
+def atan2(col1: ColumnOrName, col2: SupportsFloat) -> Column: ...
 def avg(col: ColumnOrName) -> Column: ...
 def base64(col: ColumnOrName) -> Column: ...
 def bitwiseNOT(col: ColumnOrName) -> Column: ...
@@ -186,7 +191,12 @@ def desc_nulls_last(col: ColumnOrName) -> Column: ...
 def exp(col: ColumnOrName) -> Column: ...
 def expm1(col: ColumnOrName) -> Column: ...
 def floor(col: ColumnOrName) -> Column: ...
-def hypot(col: ColumnOrName) -> Column: ...
+@overload
+def hypot(col1: ColumnOrName, col2: ColumnOrName) -> Column: ...
+@overload
+def hypot(col1: SupportsFloat, col2: ColumnOrName) -> Column: ...
+@overload
+def hypot(col1: ColumnOrName, col2: SupportsFloat) -> Column: ...
 def kurtosis(col: ColumnOrName) -> Column: ...
 def lit(col: Any) -> Column: ...
 def log10(col: ColumnOrName) -> Column: ...
@@ -197,7 +207,12 @@ def max(col: ColumnOrName) -> Column: ...
 def mean(col: ColumnOrName) -> Column: ...
 def min(col: ColumnOrName) -> Column: ...
 def percent_rank() -> Column: ...
+@overload
 def pow(col1: ColumnOrName, col2: ColumnOrName) -> Column: ...
+@overload
+def pow(col1: SupportsFloat, col2: ColumnOrName) -> Column: ...
+@overload
+def pow(col1: ColumnOrName, col2: SupportsFloat) -> Column: ...
 def radians(col: ColumnOrName) -> Column: ...
 def rank() -> Column: ...
 def rint(col: ColumnOrName) -> Column: ...