This is an automated email from the ASF dual-hosted git repository.
gurwls223 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 7aee290 [SPARK-36880][PYTHON] Inline type hints for
python/pyspark/sql/functions.py
7aee290 is described below
commit 7aee29016a3b2986846195c45cf6dea44db61738
Author: Xinrong Meng <[email protected]>
AuthorDate: Tue Oct 5 17:13:16 2021 +0900
[SPARK-36880][PYTHON] Inline type hints for python/pyspark/sql/functions.py
### What changes were proposed in this pull request?
Inline type hints from `python/pyspark/sql/functions.pyi` to
`python/pyspark/sql/functions.py`.
### Why are the changes needed?
Currently, there is type hint stub files `python/pyspark/sql/functions.pyi`
to show the expected types for functions, but we can also take advantage of
static type checking within the functions by inlining the type hints.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Existing test.
Closes #34130 from xinrong-databricks/inline_functions.
Authored-by: Xinrong Meng <[email protected]>
Signed-off-by: Hyukjin Kwon <[email protected]>
---
python/pyspark/sql/functions.py | 1043 +++++++++++++++++++++++---------------
python/pyspark/sql/functions.pyi | 377 --------------
2 files changed, 624 insertions(+), 796 deletions(-)
diff --git a/python/pyspark/sql/functions.py b/python/pyspark/sql/functions.py
index b42fbf0..4637830 100644
--- a/python/pyspark/sql/functions.py
+++ b/python/pyspark/sql/functions.py
@@ -18,22 +18,47 @@
"""
A collections of builtin functions
"""
+import inspect
import sys
import functools
import warnings
+from typing import (
+ Any,
+ Callable,
+ Dict,
+ List,
+ overload,
+ Optional,
+ Tuple,
+ TYPE_CHECKING,
+ Union,
+ ValuesView
+)
from pyspark import since, SparkContext
from pyspark.rdd import PythonEvalType
-from pyspark.sql.column import Column, _to_java_column, _to_seq,
_create_column_from_literal
+from pyspark.sql.column import ( # type: ignore[attr-defined]
+ Column,
+ _to_java_column,
+ _to_seq,
+ _create_column_from_literal
+)
from pyspark.sql.dataframe import DataFrame
-from pyspark.sql.types import StringType, DataType
+from pyspark.sql.types import ArrayType, DataType, StringType, StructType
# Keep UserDefinedFunction import for backwards compatible import; moved in
SPARK-22409
-from pyspark.sql.udf import UserDefinedFunction, _create_udf # noqa: F401
-from pyspark.sql.udf import _create_udf
+from pyspark.sql.udf import UserDefinedFunction, _create_udf # type:
ignore[attr-defined] # noqa: F401
# Keep pandas_udf and PandasUDFType import for backwards compatible import;
moved in SPARK-28264
from pyspark.sql.pandas.functions import pandas_udf, PandasUDFType # noqa:
F401
from pyspark.sql.utils import to_str
+if TYPE_CHECKING:
+ from pyspark.sql._typing import (
+ ColumnOrName,
+ DataTypeOrString,
+ UserDefinedFunctionLike,
+ )
+
+
# Note to developers: all of PySpark functions here take string as column
names whenever possible.
# Namely, if columns are referred as arguments, they can be always both Column
or string,
# even though there might be few exceptions for legacy or inevitable reasons.
@@ -41,24 +66,27 @@ from pyspark.sql.utils import to_str
# since it requires to make every single overridden definition.
-def _get_get_jvm_function(name, sc):
+def _get_get_jvm_function(name: str, sc: SparkContext) -> Callable:
"""
Retrieves JVM function identified by name from
Java gateway associated with sc.
"""
- return getattr(sc._jvm.functions, name)
+ return getattr(sc._jvm.functions, name) # type: ignore[attr-defined]
-def _invoke_function(name, *args):
+def _invoke_function(name: str, *args: Any) -> Column:
"""
Invokes JVM function identified by name with args
and wraps the result with :class:`~pyspark.sql.Column`.
"""
- jf = _get_get_jvm_function(name, SparkContext._active_spark_context)
+ jf = _get_get_jvm_function(
+ name,
+ SparkContext._active_spark_context # type: ignore[attr-defined]
+ )
return Column(jf(*args))
-def _invoke_function_over_column(name, col):
+def _invoke_function_over_column(name: str, col: "ColumnOrName") -> Column:
"""
Invokes unary JVM function identified by name
and wraps the result with :class:`~pyspark.sql.Column`.
@@ -66,7 +94,7 @@ def _invoke_function_over_column(name, col):
return _invoke_function(name, _to_java_column(col))
-def _invoke_binary_math_function(name, col1, col2):
+def _invoke_binary_math_function(name: str, col1: Any, col2: Any) -> Column:
"""
Invokes binary JVM math function identified by name
and wraps the result with :class:`~pyspark.sql.Column`.
@@ -80,13 +108,13 @@ def _invoke_binary_math_function(name, col1, col2):
)
-def _options_to_str(options=None):
+def _options_to_str(options: Optional[Dict[str, Any]] = None) -> Dict[str,
str]:
if options:
return {key: to_str(value) for (key, value) in options.items()}
return {}
-def lit(col):
+def lit(col: Any) -> Column:
"""
Creates a :class:`~pyspark.sql.Column` of literal value.
@@ -101,7 +129,7 @@ def lit(col):
@since(1.3)
-def col(col):
+def col(col: str) -> Column:
"""
Returns a :class:`~pyspark.sql.Column` based on the given column name.'
Examples
@@ -118,7 +146,7 @@ column = col
@since(1.3)
-def asc(col):
+def asc(col: "ColumnOrName") -> Column:
"""
Returns a sort expression based on the ascending order of the given column
name.
"""
@@ -129,7 +157,7 @@ def asc(col):
@since(1.3)
-def desc(col):
+def desc(col: "ColumnOrName") -> Column:
"""
Returns a sort expression based on the descending order of the given
column name.
"""
@@ -140,7 +168,7 @@ def desc(col):
@since(1.3)
-def sqrt(col):
+def sqrt(col: "ColumnOrName") -> Column:
"""
Computes the square root of the specified float value.
"""
@@ -148,7 +176,7 @@ def sqrt(col):
@since(1.3)
-def abs(col):
+def abs(col: "ColumnOrName") -> Column:
"""
Computes the absolute value.
"""
@@ -156,7 +184,7 @@ def abs(col):
@since(1.3)
-def max(col):
+def max(col: "ColumnOrName") -> Column:
"""
Aggregate function: returns the maximum value of the expression in a group.
"""
@@ -164,7 +192,7 @@ def max(col):
@since(1.3)
-def min(col):
+def min(col: "ColumnOrName") -> Column:
"""
Aggregate function: returns the minimum value of the expression in a group.
"""
@@ -172,7 +200,7 @@ def min(col):
@since(1.3)
-def count(col):
+def count(col: "ColumnOrName") -> Column:
"""
Aggregate function: returns the number of items in a group.
"""
@@ -180,7 +208,7 @@ def count(col):
@since(1.3)
-def sum(col):
+def sum(col: "ColumnOrName") -> Column:
"""
Aggregate function: returns the sum of all values in the expression.
"""
@@ -188,7 +216,7 @@ def sum(col):
@since(1.3)
-def avg(col):
+def avg(col: "ColumnOrName") -> Column:
"""
Aggregate function: returns the average of the values in a group.
"""
@@ -196,7 +224,7 @@ def avg(col):
@since(1.3)
-def mean(col):
+def mean(col: "ColumnOrName") -> Column:
"""
Aggregate function: returns the average of the values in a group.
"""
@@ -204,7 +232,7 @@ def mean(col):
@since(1.3)
-def sumDistinct(col):
+def sumDistinct(col: "ColumnOrName") -> Column:
"""
Aggregate function: returns the sum of distinct values in the expression.
@@ -216,14 +244,14 @@ def sumDistinct(col):
@since(3.2)
-def sum_distinct(col):
+def sum_distinct(col: "ColumnOrName") -> Column:
"""
Aggregate function: returns the sum of distinct values in the expression.
"""
return _invoke_function_over_column("sum_distinct", col)
-def product(col):
+def product(col: "ColumnOrName") -> Column:
"""
Aggregate function: returns the product of the values in a group.
@@ -251,7 +279,7 @@ def product(col):
return _invoke_function_over_column("product", col)
-def acos(col):
+def acos(col: "ColumnOrName") -> Column:
"""
Computes inverse cosine of the input column.
@@ -265,7 +293,7 @@ def acos(col):
return _invoke_function_over_column("acos", col)
-def acosh(col):
+def acosh(col: "ColumnOrName") -> Column:
"""
Computes inverse hyperbolic cosine of the input column.
@@ -278,7 +306,7 @@ def acosh(col):
return _invoke_function_over_column("acosh", col)
-def asin(col):
+def asin(col: "ColumnOrName") -> Column:
"""
Computes inverse sine of the input column.
@@ -293,7 +321,7 @@ def asin(col):
return _invoke_function_over_column("asin", col)
-def asinh(col):
+def asinh(col: "ColumnOrName") -> Column:
"""
Computes inverse hyperbolic sine of the input column.
@@ -306,7 +334,7 @@ def asinh(col):
return _invoke_function_over_column("asinh", col)
-def atan(col):
+def atan(col: "ColumnOrName") -> Column:
"""
Compute inverse tangent of the input column.
@@ -320,7 +348,7 @@ def atan(col):
return _invoke_function_over_column("atan", col)
-def atanh(col):
+def atanh(col: "ColumnOrName") -> Column:
"""
Computes inverse hyperbolic tangent of the input column.
@@ -334,7 +362,7 @@ def atanh(col):
@since(1.4)
-def cbrt(col):
+def cbrt(col: "ColumnOrName") -> Column:
"""
Computes the cube-root of the given value.
"""
@@ -342,14 +370,14 @@ def cbrt(col):
@since(1.4)
-def ceil(col):
+def ceil(col: "ColumnOrName") -> Column:
"""
Computes the ceiling of the given value.
"""
return _invoke_function_over_column("ceil", col)
-def cos(col):
+def cos(col: "ColumnOrName") -> Column:
"""
Computes cosine of the input column.
@@ -368,7 +396,7 @@ def cos(col):
return _invoke_function_over_column("cos", col)
-def cosh(col):
+def cosh(col: "ColumnOrName") -> Column:
"""
Computes hyperbolic cosine of the input column.
@@ -387,7 +415,7 @@ def cosh(col):
return _invoke_function_over_column("cosh", col)
-def cot(col):
+def cot(col: "ColumnOrName") -> Column:
"""
Computes cotangent of the input column.
@@ -406,7 +434,7 @@ def cot(col):
return _invoke_function_over_column("cot", col)
-def csc(col):
+def csc(col: "ColumnOrName") -> Column:
"""
Computes cosecant of the input column.
@@ -426,7 +454,7 @@ def csc(col):
@since(1.4)
-def exp(col):
+def exp(col: "ColumnOrName") -> Column:
"""
Computes the exponential of the given value.
"""
@@ -434,7 +462,7 @@ def exp(col):
@since(1.4)
-def expm1(col):
+def expm1(col: "ColumnOrName") -> Column:
"""
Computes the exponential of the given value minus one.
"""
@@ -442,7 +470,7 @@ def expm1(col):
@since(1.4)
-def floor(col):
+def floor(col: "ColumnOrName") -> Column:
"""
Computes the floor of the given value.
"""
@@ -450,7 +478,7 @@ def floor(col):
@since(1.4)
-def log(col):
+def log(col: "ColumnOrName") -> Column:
"""
Computes the natural logarithm of the given value.
"""
@@ -458,7 +486,7 @@ def log(col):
@since(1.4)
-def log10(col):
+def log10(col: "ColumnOrName") -> Column:
"""
Computes the logarithm of the given value in Base 10.
"""
@@ -466,7 +494,7 @@ def log10(col):
@since(1.4)
-def log1p(col):
+def log1p(col: "ColumnOrName") -> Column:
"""
Computes the natural logarithm of the given value plus one.
"""
@@ -474,7 +502,7 @@ def log1p(col):
@since(1.4)
-def rint(col):
+def rint(col: "ColumnOrName") -> Column:
"""
Returns the double value that is closest in value to the argument and
is equal to a mathematical integer.
@@ -482,7 +510,7 @@ def rint(col):
return _invoke_function_over_column("rint", col)
-def sec(col):
+def sec(col: "ColumnOrName") -> Column:
"""
Computes secant of the input column.
@@ -502,14 +530,14 @@ def sec(col):
@since(1.4)
-def signum(col):
+def signum(col: "ColumnOrName") -> Column:
"""
Computes the signum of the given value.
"""
return _invoke_function_over_column("signum", col)
-def sin(col):
+def sin(col: "ColumnOrName") -> Column:
"""
Computes sine of the input column.
@@ -527,7 +555,7 @@ def sin(col):
return _invoke_function_over_column("sin", col)
-def sinh(col):
+def sinh(col: "ColumnOrName") -> Column:
"""
Computes hyperbolic sine of the input column.
@@ -547,7 +575,7 @@ def sinh(col):
return _invoke_function_over_column("sinh", col)
-def tan(col):
+def tan(col: "ColumnOrName") -> Column:
"""
Computes tangent of the input column.
@@ -566,7 +594,7 @@ def tan(col):
return _invoke_function_over_column("tan", col)
-def tanh(col):
+def tanh(col: "ColumnOrName") -> Column:
"""
Computes hyperbolic tangent of the input column.
@@ -587,7 +615,7 @@ def tanh(col):
@since(1.4)
-def toDegrees(col):
+def toDegrees(col: "ColumnOrName") -> Column:
"""
.. deprecated:: 2.1.0
Use :func:`degrees` instead.
@@ -597,7 +625,7 @@ def toDegrees(col):
@since(1.4)
-def toRadians(col):
+def toRadians(col: "ColumnOrName") -> Column:
"""
.. deprecated:: 2.1.0
Use :func:`radians` instead.
@@ -607,7 +635,7 @@ def toRadians(col):
@since(1.4)
-def bitwiseNOT(col):
+def bitwiseNOT(col: "ColumnOrName") -> Column:
"""
Computes bitwise not.
@@ -619,7 +647,7 @@ def bitwiseNOT(col):
@since(3.2)
-def bitwise_not(col):
+def bitwise_not(col: "ColumnOrName") -> Column:
"""
Computes bitwise not.
"""
@@ -627,7 +655,7 @@ def bitwise_not(col):
@since(2.4)
-def asc_nulls_first(col):
+def asc_nulls_first(col: "ColumnOrName") -> Column:
"""
Returns a sort expression based on the ascending order of the given
column name, and null values return before non-null values.
@@ -639,7 +667,7 @@ def asc_nulls_first(col):
@since(2.4)
-def asc_nulls_last(col):
+def asc_nulls_last(col: "ColumnOrName") -> Column:
"""
Returns a sort expression based on the ascending order of the given
column name, and null values appear after non-null values.
@@ -651,7 +679,7 @@ def asc_nulls_last(col):
@since(2.4)
-def desc_nulls_first(col):
+def desc_nulls_first(col: "ColumnOrName") -> Column:
"""
Returns a sort expression based on the descending order of the given
column name, and null values appear before non-null values.
@@ -663,7 +691,7 @@ def desc_nulls_first(col):
@since(2.4)
-def desc_nulls_last(col):
+def desc_nulls_last(col: "ColumnOrName") -> Column:
"""
Returns a sort expression based on the descending order of the given
column name, and null values appear after non-null values.
@@ -675,7 +703,7 @@ def desc_nulls_last(col):
@since(1.6)
-def stddev(col):
+def stddev(col: "ColumnOrName") -> Column:
"""
Aggregate function: alias for stddev_samp.
"""
@@ -683,7 +711,7 @@ def stddev(col):
@since(1.6)
-def stddev_samp(col):
+def stddev_samp(col: "ColumnOrName") -> Column:
"""
Aggregate function: returns the unbiased sample standard deviation of
the expression in a group.
@@ -692,7 +720,7 @@ def stddev_samp(col):
@since(1.6)
-def stddev_pop(col):
+def stddev_pop(col: "ColumnOrName") -> Column:
"""
Aggregate function: returns population standard deviation of
the expression in a group.
@@ -701,7 +729,7 @@ def stddev_pop(col):
@since(1.6)
-def variance(col):
+def variance(col: "ColumnOrName") -> Column:
"""
Aggregate function: alias for var_samp
"""
@@ -709,7 +737,7 @@ def variance(col):
@since(1.6)
-def var_samp(col):
+def var_samp(col: "ColumnOrName") -> Column:
"""
Aggregate function: returns the unbiased sample variance of
the values in a group.
@@ -718,7 +746,7 @@ def var_samp(col):
@since(1.6)
-def var_pop(col):
+def var_pop(col: "ColumnOrName") -> Column:
"""
Aggregate function: returns the population variance of the values in a
group.
"""
@@ -726,7 +754,7 @@ def var_pop(col):
@since(1.6)
-def skewness(col):
+def skewness(col: "ColumnOrName") -> Column:
"""
Aggregate function: returns the skewness of the values in a group.
"""
@@ -734,14 +762,14 @@ def skewness(col):
@since(1.6)
-def kurtosis(col):
+def kurtosis(col: "ColumnOrName") -> Column:
"""
Aggregate function: returns the kurtosis of the values in a group.
"""
return _invoke_function_over_column("kurtosis", col)
-def collect_list(col):
+def collect_list(col: "ColumnOrName") -> Column:
"""
Aggregate function: returns a list of objects with duplicates.
@@ -761,7 +789,7 @@ def collect_list(col):
return _invoke_function_over_column("collect_list", col)
-def collect_set(col):
+def collect_set(col: "ColumnOrName") -> Column:
"""
Aggregate function: returns a set of objects with duplicate elements
eliminated.
@@ -781,7 +809,7 @@ def collect_set(col):
return _invoke_function_over_column("collect_set", col)
-def degrees(col):
+def degrees(col: "ColumnOrName") -> Column:
"""
Converts an angle measured in radians to an approximately equivalent angle
measured in degrees.
@@ -801,7 +829,7 @@ def degrees(col):
return _invoke_function_over_column("degrees", col)
-def radians(col):
+def radians(col: "ColumnOrName") -> Column:
"""
Converts an angle measured in degrees to an approximately equivalent angle
measured in radians.
@@ -821,7 +849,22 @@ def radians(col):
return _invoke_function_over_column("radians", col)
-def atan2(col1, col2):
+@overload
+def atan2(col1: "ColumnOrName", col2: "ColumnOrName") -> Column:
+ ...
+
+
+@overload
+def atan2(col1: float, col2: "ColumnOrName") -> Column:
+ ...
+
+
+@overload
+def atan2(col1: "ColumnOrName", col2: float) -> Column:
+ ...
+
+
+def atan2(col1: Union["ColumnOrName", float], col2: Union["ColumnOrName",
float]) -> Column:
"""
.. versionadded:: 1.4.0
@@ -844,16 +887,46 @@ def atan2(col1, col2):
return _invoke_binary_math_function("atan2", col1, col2)
+@overload
+def hypot(col1: "ColumnOrName", col2: "ColumnOrName") -> Column:
+ ...
+
+
+@overload
+def hypot(col1: float, col2: "ColumnOrName") -> Column:
+ ...
+
+
+@overload
+def hypot(col1: "ColumnOrName", col2: float) -> Column:
+ ...
+
+
@since(1.4)
-def hypot(col1, col2):
+def hypot(col1: Union["ColumnOrName", float], col2: Union["ColumnOrName",
float]) -> Column:
"""
Computes ``sqrt(a^2 + b^2)`` without intermediate overflow or underflow.
"""
return _invoke_binary_math_function("hypot", col1, col2)
+@overload
+def pow(col1: "ColumnOrName", col2: "ColumnOrName") -> Column:
+ ...
+
+
+@overload
+def pow(col1: float, col2: "ColumnOrName") -> Column:
+ ...
+
+
+@overload
+def pow(col1: "ColumnOrName", col2: float) -> Column:
+ ...
+
+
@since(1.4)
-def pow(col1, col2):
+def pow(col1: Union["ColumnOrName", float], col2: Union["ColumnOrName",
float]) -> Column:
"""
Returns the value of the first argument raised to the power of the second
argument.
"""
@@ -861,7 +934,7 @@ def pow(col1, col2):
@since(1.6)
-def row_number():
+def row_number() -> Column:
"""
Window function: returns a sequential number starting at 1 within a window
partition.
"""
@@ -869,7 +942,7 @@ def row_number():
@since(1.6)
-def dense_rank():
+def dense_rank() -> Column:
"""
Window function: returns the rank of rows within a window partition,
without any gaps.
@@ -885,7 +958,7 @@ def dense_rank():
@since(1.6)
-def rank():
+def rank() -> Column:
"""
Window function: returns the rank of rows within a window partition.
@@ -901,7 +974,7 @@ def rank():
@since(1.6)
-def cume_dist():
+def cume_dist() -> Column:
"""
Window function: returns the cumulative distribution of values within a
window partition,
i.e. the fraction of rows that are below the current row.
@@ -910,7 +983,7 @@ def cume_dist():
@since(1.6)
-def percent_rank():
+def percent_rank() -> Column:
"""
Window function: returns the relative rank (i.e. percentile) of rows
within a window partition.
"""
@@ -918,7 +991,7 @@ def percent_rank():
@since(1.3)
-def approxCountDistinct(col, rsd=None):
+def approxCountDistinct(col: "ColumnOrName", rsd: Optional[float] = None) ->
Column:
"""
.. deprecated:: 2.1.0
Use :func:`approx_count_distinct` instead.
@@ -927,7 +1000,7 @@ def approxCountDistinct(col, rsd=None):
return approx_count_distinct(col, rsd)
-def approx_count_distinct(col, rsd=None):
+def approx_count_distinct(col: "ColumnOrName", rsd: Optional[float] = None) ->
Column:
"""Aggregate function: returns a new :class:`~pyspark.sql.Column` for
approximate distinct count
of column `col`.
@@ -945,7 +1018,7 @@ def approx_count_distinct(col, rsd=None):
>>> df.agg(approx_count_distinct(df.age).alias('distinct_ages')).collect()
[Row(distinct_ages=2)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
if rsd is None:
jc = sc._jvm.functions.approx_count_distinct(_to_java_column(col))
else:
@@ -954,14 +1027,14 @@ def approx_count_distinct(col, rsd=None):
@since(1.6)
-def broadcast(df):
+def broadcast(df: DataFrame) -> DataFrame:
"""Marks a DataFrame as small enough for use in broadcast joins."""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return DataFrame(sc._jvm.functions.broadcast(df._jdf), df.sql_ctx)
-def coalesce(*cols):
+def coalesce(*cols: "ColumnOrName") -> Column:
"""Returns the first column that is not null.
.. versionadded:: 1.4.0
@@ -996,12 +1069,12 @@ def coalesce(*cols):
|null| 2| 0.0|
+----+----+----------------+
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
jc = sc._jvm.functions.coalesce(_to_seq(sc, cols, _to_java_column))
return Column(jc)
-def corr(col1, col2):
+def corr(col1: "ColumnOrName", col2: "ColumnOrName") -> Column:
"""Returns a new :class:`~pyspark.sql.Column` for the Pearson Correlation
Coefficient for
``col1`` and ``col2``.
@@ -1015,11 +1088,11 @@ def corr(col1, col2):
>>> df.agg(corr("a", "b").alias('c')).collect()
[Row(c=1.0)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.corr(_to_java_column(col1),
_to_java_column(col2)))
-def covar_pop(col1, col2):
+def covar_pop(col1: "ColumnOrName", col2: "ColumnOrName") -> Column:
"""Returns a new :class:`~pyspark.sql.Column` for the population
covariance of ``col1`` and
``col2``.
@@ -1033,11 +1106,11 @@ def covar_pop(col1, col2):
>>> df.agg(covar_pop("a", "b").alias('c')).collect()
[Row(c=0.0)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.covar_pop(_to_java_column(col1),
_to_java_column(col2)))
-def covar_samp(col1, col2):
+def covar_samp(col1: "ColumnOrName", col2: "ColumnOrName") -> Column:
"""Returns a new :class:`~pyspark.sql.Column` for the sample covariance of
``col1`` and
``col2``.
@@ -1051,11 +1124,11 @@ def covar_samp(col1, col2):
>>> df.agg(covar_samp("a", "b").alias('c')).collect()
[Row(c=0.0)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.covar_samp(_to_java_column(col1),
_to_java_column(col2)))
-def countDistinct(col, *cols):
+def countDistinct(col: "ColumnOrName", *cols: "ColumnOrName") -> Column:
"""Returns a new :class:`~pyspark.sql.Column` for distinct count of
``col`` or ``cols``.
An alias of :func:`count_distinct`, and it is encouraged to use
:func:`count_distinct`
@@ -1066,7 +1139,7 @@ def countDistinct(col, *cols):
return count_distinct(col, *cols)
-def count_distinct(col, *cols):
+def count_distinct(col: "ColumnOrName", *cols: "ColumnOrName") -> Column:
"""Returns a new :class:`Column` for distinct count of ``col`` or ``cols``.
.. versionadded:: 3.2.0
@@ -1079,12 +1152,12 @@ def count_distinct(col, *cols):
>>> df.agg(count_distinct("age", "name").alias('c')).collect()
[Row(c=2)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
jc = sc._jvm.functions.count_distinct(_to_java_column(col), _to_seq(sc,
cols, _to_java_column))
return Column(jc)
-def first(col, ignorenulls=False):
+def first(col: "ColumnOrName", ignorenulls: bool = False) -> Column:
"""Aggregate function: returns the first value in a group.
The function by default returns the first values it sees. It will return
the first non-null
@@ -1097,12 +1170,12 @@ def first(col, ignorenulls=False):
The function is non-deterministic because its results depends on the order
of the
rows which may be non-deterministic after a shuffle.
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
jc = sc._jvm.functions.first(_to_java_column(col), ignorenulls)
return Column(jc)
-def grouping(col):
+def grouping(col: "ColumnOrName") -> Column:
"""
Aggregate function: indicates whether a specified column in a GROUP BY
list is aggregated
or not, returns 1 for aggregated or 0 for not aggregated in the result set.
@@ -1120,12 +1193,12 @@ def grouping(col):
| Bob| 0| 5|
+-----+--------------+--------+
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
jc = sc._jvm.functions.grouping(_to_java_column(col))
return Column(jc)
-def grouping_id(*cols):
+def grouping_id(*cols: "ColumnOrName") -> Column:
"""
Aggregate function: returns the level of grouping, equals to
@@ -1149,20 +1222,20 @@ def grouping_id(*cols):
| Bob| 0| 5|
+-----+-------------+--------+
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
jc = sc._jvm.functions.grouping_id(_to_seq(sc, cols, _to_java_column))
return Column(jc)
@since(1.6)
-def input_file_name():
+def input_file_name() -> Column:
"""Creates a string column for the file name of the current Spark task.
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.input_file_name())
-def isnan(col):
+def isnan(col: "ColumnOrName") -> Column:
"""An expression that returns true iff the column is NaN.
.. versionadded:: 1.6.0
@@ -1173,11 +1246,11 @@ def isnan(col):
>>> df.select(isnan("a").alias("r1"), isnan(df.a).alias("r2")).collect()
[Row(r1=False, r2=False), Row(r1=True, r2=True)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.isnan(_to_java_column(col)))
-def isnull(col):
+def isnull(col: "ColumnOrName") -> Column:
"""An expression that returns true iff the column is null.
.. versionadded:: 1.6.0
@@ -1188,11 +1261,11 @@ def isnull(col):
>>> df.select(isnull("a").alias("r1"), isnull(df.a).alias("r2")).collect()
[Row(r1=False, r2=False), Row(r1=True, r2=True)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.isnull(_to_java_column(col)))
-def last(col, ignorenulls=False):
+def last(col: "ColumnOrName", ignorenulls: bool = False) -> Column:
"""Aggregate function: returns the last value in a group.
The function by default returns the last values it sees. It will return
the last non-null
@@ -1205,12 +1278,12 @@ def last(col, ignorenulls=False):
The function is non-deterministic because its results depends on the order
of the
rows which may be non-deterministic after a shuffle.
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
jc = sc._jvm.functions.last(_to_java_column(col), ignorenulls)
return Column(jc)
-def monotonically_increasing_id():
+def monotonically_increasing_id() -> Column:
"""A column that generates monotonically increasing 64-bit integers.
The generated ID is guaranteed to be monotonically increasing and unique,
but not consecutive.
@@ -1232,11 +1305,11 @@ def monotonically_increasing_id():
>>> df0.select(monotonically_increasing_id().alias('id')).collect()
[Row(id=0), Row(id=1), Row(id=2), Row(id=8589934592), Row(id=8589934593),
Row(id=8589934594)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.monotonically_increasing_id())
-def nanvl(col1, col2):
+def nanvl(col1: "ColumnOrName", col2: "ColumnOrName") -> Column:
"""Returns col1 if it is not NaN, or col2 if col1 is NaN.
Both inputs should be floating point columns (:class:`DoubleType` or
:class:`FloatType`).
@@ -1249,11 +1322,15 @@ def nanvl(col1, col2):
>>> df.select(nanvl("a", "b").alias("r1"), nanvl(df.a,
df.b).alias("r2")).collect()
[Row(r1=1.0, r2=1.0), Row(r1=2.0, r2=2.0)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.nanvl(_to_java_column(col1),
_to_java_column(col2)))
-def percentile_approx(col, percentage, accuracy=10000):
+def percentile_approx(
+ col: "ColumnOrName",
+ percentage: Union[Column, float, List[float], Tuple[float]],
+ accuracy: Union[Column, float] = 10000,
+) -> Column:
"""Returns the approximate `percentile` of the numeric column `col` which
is the smallest value
in the ordered `col` values (sorted from least to greatest) such that no
more than `percentage`
of `col` values is less than the value or equal to that value.
@@ -1289,7 +1366,7 @@ def percentile_approx(col, percentage, accuracy=10000):
|-- key: long (nullable = true)
|-- median: double (nullable = true)
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
if isinstance(percentage, (list, tuple)):
# A local list
@@ -1311,7 +1388,7 @@ def percentile_approx(col, percentage, accuracy=10000):
return Column(sc._jvm.functions.percentile_approx(_to_java_column(col),
percentage, accuracy))
-def rand(seed=None):
+def rand(seed: Optional[int] = None) -> Column:
"""Generates a random column with independent and identically distributed
(i.i.d.) samples
uniformly distributed in [0.0, 1.0).
@@ -1327,7 +1404,7 @@ def rand(seed=None):
[Row(age=2, name='Alice', rand=2.4052597283576684),
Row(age=5, name='Bob', rand=2.3913904055683974)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
if seed is not None:
jc = sc._jvm.functions.rand(seed)
else:
@@ -1335,7 +1412,7 @@ def rand(seed=None):
return Column(jc)
-def randn(seed=None):
+def randn(seed: Optional[int] = None) -> Column:
"""Generates a column with independent and identically distributed
(i.i.d.) samples from
the standard normal distribution.
@@ -1351,7 +1428,7 @@ def randn(seed=None):
[Row(age=2, name='Alice', randn=1.1027054481455365),
Row(age=5, name='Bob', randn=0.7400395449950132)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
if seed is not None:
jc = sc._jvm.functions.randn(seed)
else:
@@ -1359,7 +1436,7 @@ def randn(seed=None):
return Column(jc)
-def round(col, scale=0):
+def round(col: "ColumnOrName", scale: int = 0) -> Column:
"""
Round the given value to `scale` decimal places using HALF_UP rounding
mode if `scale` >= 0
or at integral part when `scale` < 0.
@@ -1371,11 +1448,11 @@ def round(col, scale=0):
>>> spark.createDataFrame([(2.5,)], ['a']).select(round('a',
0).alias('r')).collect()
[Row(r=3.0)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.round(_to_java_column(col), scale))
-def bround(col, scale=0):
+def bround(col: "ColumnOrName", scale: int = 0) -> Column:
"""
Round the given value to `scale` decimal places using HALF_EVEN rounding
mode if `scale` >= 0
or at integral part when `scale` < 0.
@@ -1387,11 +1464,11 @@ def bround(col, scale=0):
>>> spark.createDataFrame([(2.5,)], ['a']).select(bround('a',
0).alias('r')).collect()
[Row(r=2.0)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.bround(_to_java_column(col), scale))
-def shiftLeft(col, numBits):
+def shiftLeft(col: "ColumnOrName", numBits: int) -> Column:
"""Shift the given value numBits left.
.. versionadded:: 1.5.0
@@ -1403,7 +1480,7 @@ def shiftLeft(col, numBits):
return shiftleft(col, numBits)
-def shiftleft(col, numBits):
+def shiftleft(col: "ColumnOrName", numBits: int) -> Column:
"""Shift the given value numBits left.
.. versionadded:: 3.2.0
@@ -1413,11 +1490,11 @@ def shiftleft(col, numBits):
>>> spark.createDataFrame([(21,)], ['a']).select(shiftleft('a',
1).alias('r')).collect()
[Row(r=42)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.shiftleft(_to_java_column(col), numBits))
-def shiftRight(col, numBits):
+def shiftRight(col: "ColumnOrName", numBits: int) -> Column:
"""(Signed) shift the given value numBits right.
.. versionadded:: 1.5.0
@@ -1429,7 +1506,7 @@ def shiftRight(col, numBits):
return shiftright(col, numBits)
-def shiftright(col, numBits):
+def shiftright(col: "ColumnOrName", numBits: int) -> Column:
"""(Signed) shift the given value numBits right.
.. versionadded:: 3.2.0
@@ -1439,12 +1516,12 @@ def shiftright(col, numBits):
>>> spark.createDataFrame([(42,)], ['a']).select(shiftright('a',
1).alias('r')).collect()
[Row(r=21)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
jc = sc._jvm.functions.shiftRight(_to_java_column(col), numBits)
return Column(jc)
-def shiftRightUnsigned(col, numBits):
+def shiftRightUnsigned(col: "ColumnOrName", numBits: int) -> Column:
"""Unsigned shift the given value numBits right.
.. versionadded:: 1.5.0
@@ -1456,7 +1533,7 @@ def shiftRightUnsigned(col, numBits):
return shiftrightunsigned(col, numBits)
-def shiftrightunsigned(col, numBits):
+def shiftrightunsigned(col: "ColumnOrName", numBits: int) -> Column:
"""Unsigned shift the given value numBits right.
.. versionadded:: 3.2.0
@@ -1467,12 +1544,12 @@ def shiftrightunsigned(col, numBits):
>>> df.select(shiftrightunsigned('a', 1).alias('r')).collect()
[Row(r=9223372036854775787)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
jc = sc._jvm.functions.shiftRightUnsigned(_to_java_column(col), numBits)
return Column(jc)
-def spark_partition_id():
+def spark_partition_id() -> Column:
"""A column for partition ID.
.. versionadded:: 1.6.0
@@ -1486,11 +1563,11 @@ def spark_partition_id():
>>> df.repartition(1).select(spark_partition_id().alias("pid")).collect()
[Row(pid=0), Row(pid=0)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.spark_partition_id())
-def expr(str):
+def expr(str: str) -> Column:
"""Parses the expression string into the column that it represents
.. versionadded:: 1.5.0
@@ -1500,11 +1577,11 @@ def expr(str):
>>> df.select(expr("length(name)")).collect()
[Row(length(name)=5), Row(length(name)=3)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.expr(str))
-def struct(*cols):
+def struct(*cols: "ColumnOrName") -> Column:
"""Creates a new struct column.
.. versionadded:: 1.4.0
@@ -1521,14 +1598,14 @@ def struct(*cols):
>>> df.select(struct([df.age, df.name]).alias("struct")).collect()
[Row(struct=Row(age=2, name='Alice')), Row(struct=Row(age=5, name='Bob'))]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
if len(cols) == 1 and isinstance(cols[0], (list, set)):
cols = cols[0]
jc = sc._jvm.functions.struct(_to_seq(sc, cols, _to_java_column))
return Column(jc)
-def greatest(*cols):
+def greatest(*cols: "ColumnOrName") -> Column:
"""
Returns the greatest value of the list of column names, skipping null
values.
This function takes at least 2 parameters. It will return null iff all
parameters are null.
@@ -1543,11 +1620,11 @@ def greatest(*cols):
"""
if len(cols) < 2:
raise ValueError("greatest should take at least two columns")
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.greatest(_to_seq(sc, cols,
_to_java_column)))
-def least(*cols):
+def least(*cols: Column) -> Column:
"""
Returns the least value of the list of column names, skipping null values.
This function takes at least 2 parameters. It will return null iff all
parameters are null.
@@ -1562,11 +1639,11 @@ def least(*cols):
"""
if len(cols) < 2:
raise ValueError("least should take at least two columns")
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.least(_to_seq(sc, cols, _to_java_column)))
-def when(condition, value):
+def when(condition: Column, value: Any) -> Column:
"""Evaluates a list of conditions and returns one of multiple possible
result expressions.
If :func:`pyspark.sql.Column.otherwise` is not invoked, None is returned
for unmatched
conditions.
@@ -1586,7 +1663,7 @@ def when(condition, value):
>>> df.select(when(df.age == 2, df.age + 1).alias("age")).collect()
[Row(age=3), Row(age=None)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
if not isinstance(condition, Column):
raise TypeError("condition should be a Column")
v = value._jc if isinstance(value, Column) else value
@@ -1594,7 +1671,17 @@ def when(condition, value):
return Column(jc)
-def log(arg1, arg2=None):
+@overload # type: ignore[no-redef]
+def log(arg1: "ColumnOrName") -> Column:
+ ...
+
+
+@overload
+def log(arg1: float, arg2: "ColumnOrName") -> Column:
+ ...
+
+
+def log(arg1: Union["ColumnOrName", float], arg2: Optional["ColumnOrName"] =
None) -> Column:
"""Returns the first argument-based logarithm of the second argument.
If there is only one argument, then this takes the natural logarithm of
the argument.
@@ -1609,7 +1696,7 @@ def log(arg1, arg2=None):
>>> df.select(log(df.age).alias('e')).rdd.map(lambda l:
str(l.e)[:7]).collect()
['0.69314', '1.60943']
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
if arg2 is None:
jc = sc._jvm.functions.log(_to_java_column(arg1))
else:
@@ -1617,7 +1704,7 @@ def log(arg1, arg2=None):
return Column(jc)
-def log2(col):
+def log2(col: "ColumnOrName") -> Column:
"""Returns the base-2 logarithm of the argument.
.. versionadded:: 1.5.0
@@ -1627,11 +1714,11 @@ def log2(col):
>>> spark.createDataFrame([(4,)],
['a']).select(log2('a').alias('log2')).collect()
[Row(log2=2.0)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.log2(_to_java_column(col)))
-def conv(col, fromBase, toBase):
+def conv(col: "ColumnOrName", fromBase: int, toBase: int) -> Column:
"""
Convert a number in a string column from one base to another.
@@ -1643,11 +1730,11 @@ def conv(col, fromBase, toBase):
>>> df.select(conv(df.n, 2, 16).alias('hex')).collect()
[Row(hex='15')]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.conv(_to_java_column(col), fromBase,
toBase))
-def factorial(col):
+def factorial(col: "ColumnOrName") -> Column:
"""
Computes the factorial of the given value.
@@ -1659,13 +1746,15 @@ def factorial(col):
>>> df.select(factorial(df.n).alias('f')).collect()
[Row(f=120)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.factorial(_to_java_column(col)))
# --------------- Window functions ------------------------
-def lag(col, offset=1, default=None):
+def lag(
+ col: "ColumnOrName", offset: int = 1, default: Optional[Any] = None
+) -> Column:
"""
Window function: returns the value that is `offset` rows before the
current row, and
`default` if there is less than `offset` rows before the current row. For
example,
@@ -1684,11 +1773,13 @@ def lag(col, offset=1, default=None):
default : optional
default value
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.lag(_to_java_column(col), offset, default))
-def lead(col, offset=1, default=None):
+def lead(
+ col: "ColumnOrName", offset: int = 1, default: Optional[Any] = None
+) -> Column:
"""
Window function: returns the value that is `offset` rows after the current
row, and
`default` if there is less than `offset` rows after the current row. For
example,
@@ -1707,11 +1798,13 @@ def lead(col, offset=1, default=None):
default : optional
default value
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.lead(_to_java_column(col), offset,
default))
-def nth_value(col, offset, ignoreNulls=False):
+def nth_value(
+ col: "ColumnOrName", offset: int, ignoreNulls: Optional[bool] = False
+) -> Column:
"""
Window function: returns the value that is the `offset`\\th row of the
window frame
(counting from 1), and `null` if the size of window frame is less than
`offset` rows.
@@ -1733,11 +1826,11 @@ def nth_value(col, offset, ignoreNulls=False):
indicates the Nth value should skip null in the
determination of which row to use
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.nth_value(_to_java_column(col), offset,
ignoreNulls))
-def ntile(n):
+def ntile(n: int) -> Column:
"""
Window function: returns the ntile group id (from 1 to `n` inclusive)
in an ordered window partition. For example, if `n` is 4, the first
@@ -1753,32 +1846,32 @@ def ntile(n):
n : int
an integer
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.ntile(int(n)))
# ---------------------- Date/Timestamp functions
------------------------------
@since(1.5)
-def current_date():
+def current_date() -> Column:
"""
Returns the current date at the start of query evaluation as a
:class:`DateType` column.
All calls of current_date within the same query return the same value.
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.current_date())
-def current_timestamp():
+def current_timestamp() -> Column:
"""
Returns the current timestamp at the start of query evaluation as a
:class:`TimestampType`
column. All calls of current_timestamp within the same query return the
same value.
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.current_timestamp())
-def date_format(date, format):
+def date_format(date: "ColumnOrName", format: str) -> Column:
"""
Converts a date/timestamp/string to a value of string in the format
specified by the date
format given by the second argument.
@@ -1800,11 +1893,11 @@ def date_format(date, format):
>>> df.select(date_format('dt', 'MM/dd/yyy').alias('date')).collect()
[Row(date='04/08/2015')]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.date_format(_to_java_column(date), format))
-def year(col):
+def year(col: "ColumnOrName") -> Column:
"""
Extract the year of a given date as integer.
@@ -1816,11 +1909,11 @@ def year(col):
>>> df.select(year('dt').alias('year')).collect()
[Row(year=2015)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.year(_to_java_column(col)))
-def quarter(col):
+def quarter(col: "ColumnOrName") -> Column:
"""
Extract the quarter of a given date as integer.
@@ -1832,11 +1925,11 @@ def quarter(col):
>>> df.select(quarter('dt').alias('quarter')).collect()
[Row(quarter=2)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.quarter(_to_java_column(col)))
-def month(col):
+def month(col: "ColumnOrName") -> Column:
"""
Extract the month of a given date as integer.
@@ -1848,11 +1941,11 @@ def month(col):
>>> df.select(month('dt').alias('month')).collect()
[Row(month=4)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.month(_to_java_column(col)))
-def dayofweek(col):
+def dayofweek(col: "ColumnOrName") -> Column:
"""
Extract the day of the week of a given date as integer.
Ranges from 1 for a Sunday through to 7 for a Saturday
@@ -1865,11 +1958,11 @@ def dayofweek(col):
>>> df.select(dayofweek('dt').alias('day')).collect()
[Row(day=4)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.dayofweek(_to_java_column(col)))
-def dayofmonth(col):
+def dayofmonth(col: "ColumnOrName") -> Column:
"""
Extract the day of the month of a given date as integer.
@@ -1881,11 +1974,11 @@ def dayofmonth(col):
>>> df.select(dayofmonth('dt').alias('day')).collect()
[Row(day=8)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.dayofmonth(_to_java_column(col)))
-def dayofyear(col):
+def dayofyear(col: "ColumnOrName") -> Column:
"""
Extract the day of the year of a given date as integer.
@@ -1897,11 +1990,11 @@ def dayofyear(col):
>>> df.select(dayofyear('dt').alias('day')).collect()
[Row(day=98)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.dayofyear(_to_java_column(col)))
-def hour(col):
+def hour(col: "ColumnOrName") -> Column:
"""
Extract the hours of a given date as integer.
@@ -1913,11 +2006,11 @@ def hour(col):
>>> df.select(hour('ts').alias('hour')).collect()
[Row(hour=13)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.hour(_to_java_column(col)))
-def minute(col):
+def minute(col: "ColumnOrName") -> Column:
"""
Extract the minutes of a given date as integer.
@@ -1929,11 +2022,11 @@ def minute(col):
>>> df.select(minute('ts').alias('minute')).collect()
[Row(minute=8)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.minute(_to_java_column(col)))
-def second(col):
+def second(col: "ColumnOrName") -> Column:
"""
Extract the seconds of a given date as integer.
@@ -1945,11 +2038,11 @@ def second(col):
>>> df.select(second('ts').alias('second')).collect()
[Row(second=15)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.second(_to_java_column(col)))
-def weekofyear(col):
+def weekofyear(col: "ColumnOrName") -> Column:
"""
Extract the week number of a given date as integer.
A week is considered to start on a Monday and week 1 is the first week
with more than 3 days,
@@ -1963,11 +2056,11 @@ def weekofyear(col):
>>> df.select(weekofyear(df.dt).alias('week')).collect()
[Row(week=15)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.weekofyear(_to_java_column(col)))
-def date_add(start, days):
+def date_add(start: "ColumnOrName", days: int) -> Column:
"""
Returns the date that is `days` days after `start`
@@ -1979,11 +2072,11 @@ def date_add(start, days):
>>> df.select(date_add(df.dt, 1).alias('next_date')).collect()
[Row(next_date=datetime.date(2015, 4, 9))]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.date_add(_to_java_column(start), days))
-def date_sub(start, days):
+def date_sub(start: "ColumnOrName", days: int) -> Column:
"""
Returns the date that is `days` days before `start`
@@ -1995,11 +2088,11 @@ def date_sub(start, days):
>>> df.select(date_sub(df.dt, 1).alias('prev_date')).collect()
[Row(prev_date=datetime.date(2015, 4, 7))]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.date_sub(_to_java_column(start), days))
-def datediff(end, start):
+def datediff(end: "ColumnOrName", start: "ColumnOrName") -> Column:
"""
Returns the number of days from `start` to `end`.
@@ -2011,11 +2104,11 @@ def datediff(end, start):
>>> df.select(datediff(df.d2, df.d1).alias('diff')).collect()
[Row(diff=32)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.datediff(_to_java_column(end),
_to_java_column(start)))
-def add_months(start, months):
+def add_months(start: "ColumnOrName", months: int) -> Column:
"""
Returns the date that is `months` months after `start`
@@ -2027,11 +2120,13 @@ def add_months(start, months):
>>> df.select(add_months(df.dt, 1).alias('next_month')).collect()
[Row(next_month=datetime.date(2015, 5, 8))]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.add_months(_to_java_column(start), months))
-def months_between(date1, date2, roundOff=True):
+def months_between(
+ date1: "ColumnOrName", date2: "ColumnOrName", roundOff: bool = True
+) -> Column:
"""
Returns number of months between dates date1 and date2.
If date1 is later than date2, then the result is positive.
@@ -2049,12 +2144,12 @@ def months_between(date1, date2, roundOff=True):
>>> df.select(months_between(df.date1, df.date2,
False).alias('months')).collect()
[Row(months=3.9495967741935485)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.months_between(
_to_java_column(date1), _to_java_column(date2), roundOff))
-def to_date(col, format=None):
+def to_date(col: "ColumnOrName", format: Optional[str] = None) -> Column:
"""Converts a :class:`~pyspark.sql.Column` into
:class:`pyspark.sql.types.DateType`
using the optionally specified format. Specify formats according to
`datetime pattern`_.
By default, it follows casting rules to
:class:`pyspark.sql.types.DateType` if the format
@@ -2074,7 +2169,7 @@ def to_date(col, format=None):
>>> df.select(to_date(df.t, 'yyyy-MM-dd HH:mm:ss').alias('date')).collect()
[Row(date=datetime.date(1997, 2, 28))]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
if format is None:
jc = sc._jvm.functions.to_date(_to_java_column(col))
else:
@@ -2082,7 +2177,17 @@ def to_date(col, format=None):
return Column(jc)
-def to_timestamp(col, format=None):
+@overload
+def to_timestamp(col: "ColumnOrName") -> Column:
+ ...
+
+
+@overload
+def to_timestamp(col: "ColumnOrName", format: str) -> Column:
+ ...
+
+
+def to_timestamp(col: "ColumnOrName", format: Optional[str] = None) -> Column:
"""Converts a :class:`~pyspark.sql.Column` into
:class:`pyspark.sql.types.TimestampType`
using the optionally specified format. Specify formats according to
`datetime pattern`_.
By default, it follows casting rules to
:class:`pyspark.sql.types.TimestampType` if the format
@@ -2102,7 +2207,7 @@ def to_timestamp(col, format=None):
>>> df.select(to_timestamp(df.t, 'yyyy-MM-dd
HH:mm:ss').alias('dt')).collect()
[Row(dt=datetime.datetime(1997, 2, 28, 10, 30))]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
if format is None:
jc = sc._jvm.functions.to_timestamp(_to_java_column(col))
else:
@@ -2110,7 +2215,7 @@ def to_timestamp(col, format=None):
return Column(jc)
-def trunc(date, format):
+def trunc(date: "ColumnOrName", format: str) -> Column:
"""
Returns date truncated to the unit specified by the format.
@@ -2132,11 +2237,11 @@ def trunc(date, format):
>>> df.select(trunc(df.d, 'mon').alias('month')).collect()
[Row(month=datetime.date(1997, 2, 1))]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.trunc(_to_java_column(date), format))
-def date_trunc(format, timestamp):
+def date_trunc(format: str, timestamp: "ColumnOrName") -> Column:
"""
Returns timestamp truncated to the unit specified by the format.
@@ -2160,11 +2265,11 @@ def date_trunc(format, timestamp):
>>> df.select(date_trunc('mon', df.t).alias('month')).collect()
[Row(month=datetime.datetime(1997, 2, 1, 0, 0))]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.date_trunc(format,
_to_java_column(timestamp)))
-def next_day(date, dayOfWeek):
+def next_day(date: "ColumnOrName", dayOfWeek: str) -> Column:
"""
Returns the first date which is later than the value of the date column.
@@ -2179,11 +2284,11 @@ def next_day(date, dayOfWeek):
>>> df.select(next_day(df.d, 'Sun').alias('date')).collect()
[Row(date=datetime.date(2015, 8, 2))]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.next_day(_to_java_column(date), dayOfWeek))
-def last_day(date):
+def last_day(date: "ColumnOrName") -> Column:
"""
Returns the last day of the month which the given date belongs to.
@@ -2195,11 +2300,11 @@ def last_day(date):
>>> df.select(last_day(df.d).alias('date')).collect()
[Row(date=datetime.date(1997, 2, 28))]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.last_day(_to_java_column(date)))
-def from_unixtime(timestamp, format="yyyy-MM-dd HH:mm:ss"):
+def from_unixtime(timestamp: "ColumnOrName", format: str = "yyyy-MM-dd
HH:mm:ss") -> Column:
"""
Converts the number of seconds from unix epoch (1970-01-01 00:00:00 UTC)
to a string
representing the timestamp of that moment in the current system time zone
in the given
@@ -2215,11 +2320,13 @@ def from_unixtime(timestamp, format="yyyy-MM-dd
HH:mm:ss"):
[Row(ts='2015-04-08 00:00:00')]
>>> spark.conf.unset("spark.sql.session.timeZone")
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.from_unixtime(_to_java_column(timestamp),
format))
-def unix_timestamp(timestamp=None, format='yyyy-MM-dd HH:mm:ss'):
+def unix_timestamp(
+ timestamp: Optional["ColumnOrName"] = None, format: str = 'yyyy-MM-dd
HH:mm:ss'
+) -> Column:
"""
Convert time string with given pattern ('yyyy-MM-dd HH:mm:ss', by default)
to Unix time stamp (in seconds), using the default timezone and the default
@@ -2237,13 +2344,13 @@ def unix_timestamp(timestamp=None, format='yyyy-MM-dd
HH:mm:ss'):
[Row(unix_time=1428476400)]
>>> spark.conf.unset("spark.sql.session.timeZone")
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
if timestamp is None:
return Column(sc._jvm.functions.unix_timestamp())
return Column(sc._jvm.functions.unix_timestamp(_to_java_column(timestamp),
format))
-def from_utc_timestamp(timestamp, tz):
+def from_utc_timestamp(timestamp: "ColumnOrName", tz: "ColumnOrName") ->
Column:
"""
This is a common function for databases supporting TIMESTAMP WITHOUT
TIMEZONE. This function
takes a timestamp which is timezone-agnostic, and interprets it as a
timestamp in UTC, and
@@ -2283,13 +2390,13 @@ def from_utc_timestamp(timestamp, tz):
>>> df.select(from_utc_timestamp(df.ts,
df.tz).alias('local_time')).collect()
[Row(local_time=datetime.datetime(1997, 2, 28, 19, 30))]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
if isinstance(tz, Column):
tz = _to_java_column(tz)
return
Column(sc._jvm.functions.from_utc_timestamp(_to_java_column(timestamp), tz))
-def to_utc_timestamp(timestamp, tz):
+def to_utc_timestamp(timestamp: "ColumnOrName", tz: "ColumnOrName") -> Column:
"""
This is a common function for databases supporting TIMESTAMP WITHOUT
TIMEZONE. This function
takes a timestamp which is timezone-agnostic, and interprets it as a
timestamp in the given
@@ -2329,13 +2436,13 @@ def to_utc_timestamp(timestamp, tz):
>>> df.select(to_utc_timestamp(df.ts, df.tz).alias('utc_time')).collect()
[Row(utc_time=datetime.datetime(1997, 2, 28, 1, 30))]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
if isinstance(tz, Column):
tz = _to_java_column(tz)
return
Column(sc._jvm.functions.to_utc_timestamp(_to_java_column(timestamp), tz))
-def timestamp_seconds(col):
+def timestamp_seconds(col: "ColumnOrName") -> Column:
"""
.. versionadded:: 3.1.0
@@ -2353,11 +2460,16 @@ def timestamp_seconds(col):
>>> spark.conf.unset("spark.sql.session.timeZone")
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.timestamp_seconds(_to_java_column(col)))
-def window(timeColumn, windowDuration, slideDuration=None, startTime=None):
+def window(
+ timeColumn: "ColumnOrName",
+ windowDuration: str,
+ slideDuration: Optional[str] = None,
+ startTime: Optional[str] = None,
+) -> Column:
"""Bucketize rows into one or more time windows given a timestamp
specifying column. Window
starts are inclusive but the window ends are exclusive, e.g. 12:05 will be
in the window
[12:05,12:10) but not in [12:00,12:05). Windows can support microsecond
precision. Windows in
@@ -2409,11 +2521,11 @@ def window(timeColumn, windowDuration,
slideDuration=None, startTime=None):
... w.window.end.cast("string").alias("end"), "sum").collect()
[Row(start='2016-03-11 09:00:05', end='2016-03-11 09:00:10', sum=1)]
"""
- def check_string_field(field, fieldName):
+ def check_string_field(field, fieldName): # type: ignore[no-untyped-def]
if not field or type(field) is not str:
raise TypeError("%s should be provided as a string" % fieldName)
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
time_col = _to_java_column(timeColumn)
check_string_field(windowDuration, "windowDuration")
if slideDuration and startTime:
@@ -2431,7 +2543,7 @@ def window(timeColumn, windowDuration,
slideDuration=None, startTime=None):
return Column(res)
-def session_window(timeColumn, gapDuration):
+def session_window(timeColumn: "ColumnOrName", gapDuration: Union[Column,
str]) -> Column:
"""
Generates session window given a timestamp specifying column.
Session window is one of dynamic windows, which means the length of window
is varying
@@ -2473,11 +2585,11 @@ def session_window(timeColumn, gapDuration):
... w.session_window.end.cast("string").alias("end"),
"sum").collect()
[Row(start='2016-03-11 09:00:07', end='2016-03-11 09:00:12', sum=1)]
"""
- def check_field(field, fieldName):
+ def check_field(field: Union[Column, str], fieldName: str) -> None:
if field is None or not isinstance(field, (str, Column)):
raise TypeError("%s should be provided as a string or Column" %
fieldName)
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
time_col = _to_java_column(timeColumn)
check_field(gapDuration, "gapDuration")
gap_duration = (
@@ -2491,7 +2603,7 @@ def session_window(timeColumn, gapDuration):
# ---------------------------- misc functions
----------------------------------
-def crc32(col):
+def crc32(col: "ColumnOrName") -> Column:
"""
Calculates the cyclic redundancy check value (CRC32) of a binary column
and
returns the value as a bigint.
@@ -2503,11 +2615,11 @@ def crc32(col):
>>> spark.createDataFrame([('ABC',)],
['a']).select(crc32('a').alias('crc32')).collect()
[Row(crc32=2743272264)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.crc32(_to_java_column(col)))
-def md5(col):
+def md5(col: "ColumnOrName") -> Column:
"""Calculates the MD5 digest and returns the value as a 32 character hex
string.
.. versionadded:: 1.5.0
@@ -2517,12 +2629,12 @@ def md5(col):
>>> spark.createDataFrame([('ABC',)],
['a']).select(md5('a').alias('hash')).collect()
[Row(hash='902fbdd2b1df0c4f70b4a5d23525e932')]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
jc = sc._jvm.functions.md5(_to_java_column(col))
return Column(jc)
-def sha1(col):
+def sha1(col: "ColumnOrName") -> Column:
"""Returns the hex string result of SHA-1.
.. versionadded:: 1.5.0
@@ -2532,12 +2644,12 @@ def sha1(col):
>>> spark.createDataFrame([('ABC',)],
['a']).select(sha1('a').alias('hash')).collect()
[Row(hash='3c01bdbb26f358bab27f267924aa2c9a03fcfdb8')]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
jc = sc._jvm.functions.sha1(_to_java_column(col))
return Column(jc)
-def sha2(col, numBits):
+def sha2(col: "ColumnOrName", numBits: int) -> Column:
"""Returns the hex string result of SHA-2 family of hash functions
(SHA-224, SHA-256, SHA-384,
and SHA-512). The numBits indicates the desired bit length of the result,
which must have a
value of 224, 256, 384, 512, or 0 (which is equivalent to 256).
@@ -2552,12 +2664,12 @@ def sha2(col, numBits):
>>> digests[1]
Row(s='cd9fb1e148ccd8442e5aa74904cc73bf6fb54d1d54d333bd596aa9bb4bb4e961')
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
jc = sc._jvm.functions.sha2(_to_java_column(col), numBits)
return Column(jc)
-def hash(*cols):
+def hash(*cols: "ColumnOrName") -> Column:
"""Calculates the hash code of given columns, and returns the result as an
int column.
.. versionadded:: 2.0.0
@@ -2567,12 +2679,12 @@ def hash(*cols):
>>> spark.createDataFrame([('ABC',)],
['a']).select(hash('a').alias('hash')).collect()
[Row(hash=-757602832)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
jc = sc._jvm.functions.hash(_to_seq(sc, cols, _to_java_column))
return Column(jc)
-def xxhash64(*cols):
+def xxhash64(*cols: "ColumnOrName") -> Column:
"""Calculates the hash code of given columns using the 64-bit variant of
the xxHash algorithm,
and returns the result as a long column.
@@ -2583,12 +2695,12 @@ def xxhash64(*cols):
>>> spark.createDataFrame([('ABC',)],
['a']).select(xxhash64('a').alias('hash')).collect()
[Row(hash=4105715581806190027)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
jc = sc._jvm.functions.xxhash64(_to_seq(sc, cols, _to_java_column))
return Column(jc)
-def assert_true(col, errMsg=None):
+def assert_true(col: "ColumnOrName", errMsg: Optional[Union[Column, str]] =
None) -> Column:
"""
Returns null if the input column is true; throws an exception with the
provided error message
otherwise.
@@ -2607,7 +2719,7 @@ def assert_true(col, errMsg=None):
>>> df.select(assert_true(df.a < df.b, 'error').alias('r')).collect()
[Row(r=None)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
if errMsg is None:
return Column(sc._jvm.functions.assert_true(_to_java_column(col)))
if not isinstance(errMsg, (str, Column)):
@@ -2624,7 +2736,7 @@ def assert_true(col, errMsg=None):
@since(3.1)
-def raise_error(errMsg):
+def raise_error(errMsg: Union[Column, str]) -> Column:
"""
Throws an exception with the provided error message.
"""
@@ -2633,7 +2745,7 @@ def raise_error(errMsg):
"errMsg should be a Column or a str, got {}".format(type(errMsg))
)
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
errMsg = (
_create_column_from_literal(errMsg)
if isinstance(errMsg, str)
@@ -2645,7 +2757,7 @@ def raise_error(errMsg):
# ---------------------- String/Binary functions ------------------------------
@since(1.5)
-def upper(col):
+def upper(col: "ColumnOrName") -> Column:
"""
Converts a string expression to upper case.
"""
@@ -2653,7 +2765,7 @@ def upper(col):
@since(1.5)
-def lower(col):
+def lower(col: "ColumnOrName") -> Column:
"""
Converts a string expression to lower case.
"""
@@ -2661,7 +2773,7 @@ def lower(col):
@since(1.5)
-def ascii(col):
+def ascii(col: "ColumnOrName") -> Column:
"""
Computes the numeric value of the first character of the string column.
"""
@@ -2669,7 +2781,7 @@ def ascii(col):
@since(1.5)
-def base64(col):
+def base64(col: "ColumnOrName") -> Column:
"""
Computes the BASE64 encoding of a binary column and returns it as a string
column.
"""
@@ -2677,7 +2789,7 @@ def base64(col):
@since(1.5)
-def unbase64(col):
+def unbase64(col: "ColumnOrName") -> Column:
"""
Decodes a BASE64 encoded string column and returns it as a binary column.
"""
@@ -2685,7 +2797,7 @@ def unbase64(col):
@since(1.5)
-def ltrim(col):
+def ltrim(col: "ColumnOrName") -> Column:
"""
Trim the spaces from left end for the specified string value.
"""
@@ -2693,7 +2805,7 @@ def ltrim(col):
@since(1.5)
-def rtrim(col):
+def rtrim(col: "ColumnOrName") -> Column:
"""
Trim the spaces from right end for the specified string value.
"""
@@ -2701,14 +2813,14 @@ def rtrim(col):
@since(1.5)
-def trim(col):
+def trim(col: "ColumnOrName") -> Column:
"""
Trim the spaces from both ends for the specified string column.
"""
return _invoke_function_over_column("trim", col)
-def concat_ws(sep, *cols):
+def concat_ws(sep: str, *cols: "ColumnOrName") -> Column:
"""
Concatenates multiple input string columns together into a single string
column,
using the given separator.
@@ -2721,31 +2833,31 @@ def concat_ws(sep, *cols):
>>> df.select(concat_ws('-', df.s, df.d).alias('s')).collect()
[Row(s='abcd-123')]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.concat_ws(sep, _to_seq(sc, cols,
_to_java_column)))
@since(1.5)
-def decode(col, charset):
+def decode(col: "ColumnOrName", charset: str) -> Column:
"""
Computes the first argument into a string from a binary using the provided
character set
(one of 'US-ASCII', 'ISO-8859-1', 'UTF-8', 'UTF-16BE', 'UTF-16LE',
'UTF-16').
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.decode(_to_java_column(col), charset))
@since(1.5)
-def encode(col, charset):
+def encode(col: "ColumnOrName", charset: str) -> Column:
"""
Computes the first argument into a binary from a string using the provided
character set
(one of 'US-ASCII', 'ISO-8859-1', 'UTF-8', 'UTF-16BE', 'UTF-16LE',
'UTF-16').
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.encode(_to_java_column(col), charset))
-def format_number(col, d):
+def format_number(col: "ColumnOrName", d: int) -> Column:
"""
Formats the number X to a format like '#,--#,--#.--', rounded to d decimal
places
with HALF_EVEN round mode, and returns the result as a string.
@@ -2762,11 +2874,11 @@ def format_number(col, d):
>>> spark.createDataFrame([(5,)], ['a']).select(format_number('a',
4).alias('v')).collect()
[Row(v='5.0000')]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.format_number(_to_java_column(col), d))
-def format_string(format, *cols):
+def format_string(format: str, *cols: "ColumnOrName") -> Column:
"""
Formats the arguments in printf-style and returns the result as a string
column.
@@ -2785,11 +2897,11 @@ def format_string(format, *cols):
>>> df.select(format_string('%d %s', df.a, df.b).alias('v')).collect()
[Row(v='5 hello')]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.format_string(format, _to_seq(sc, cols,
_to_java_column)))
-def instr(str, substr):
+def instr(str: "ColumnOrName", substr: str) -> Column:
"""
Locate the position of the first occurrence of substr column in the given
string.
Returns null if either of the arguments are null.
@@ -2805,11 +2917,16 @@ def instr(str, substr):
>>> df.select(instr(df.s, 'b').alias('s')).collect()
[Row(s=2)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.instr(_to_java_column(str), substr))
-def overlay(src, replace, pos, len=-1):
+def overlay(
+ src: "ColumnOrName",
+ replace: "ColumnOrName",
+ pos: Union[Column, int],
+ len: Union[Column, int] = -1,
+) -> Column:
"""
Overlay the specified portion of `src` with `replace`,
starting from byte position `pos` of `src` and proceeding for `len` bytes.
@@ -2836,7 +2953,7 @@ def overlay(src, replace, pos, len=-1):
pos = _create_column_from_literal(pos) if isinstance(pos, int) else
_to_java_column(pos)
len = _create_column_from_literal(len) if isinstance(len, int) else
_to_java_column(len)
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.overlay(
_to_java_column(src),
@@ -2846,7 +2963,11 @@ def overlay(src, replace, pos, len=-1):
))
-def sentences(string, language=None, country=None):
+def sentences(
+ string: "ColumnOrName",
+ language: Optional["ColumnOrName"] = None,
+ country: Optional["ColumnOrName"] = None,
+) -> Column:
"""
Splits a string into arrays of sentences, where each sentence is an array
of words.
The 'language' and 'country' arguments are optional, and if omitted, the
default locale is used.
@@ -2877,7 +2998,7 @@ def sentences(string, language=None, country=None):
if country is None:
country = lit("")
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.sentences(
_to_java_column(string),
_to_java_column(language),
@@ -2885,7 +3006,7 @@ def sentences(string, language=None, country=None):
))
-def substring(str, pos, len):
+def substring(str: "ColumnOrName", pos: int, len: int) -> Column:
"""
Substring starts at `pos` and is of length `len` when str is String type or
returns the slice of byte array that starts at `pos` in byte and is of
length `len`
@@ -2903,11 +3024,11 @@ def substring(str, pos, len):
>>> df.select(substring(df.s, 1, 2).alias('s')).collect()
[Row(s='ab')]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.substring(_to_java_column(str), pos, len))
-def substring_index(str, delim, count):
+def substring_index(str: "ColumnOrName", delim: str, count: int) -> Column:
"""
Returns the substring from string str before count occurrences of the
delimiter delim.
If count is positive, everything the left of the final delimiter (counting
from left) is
@@ -2924,11 +3045,11 @@ def substring_index(str, delim, count):
>>> df.select(substring_index(df.s, '.', -3).alias('s')).collect()
[Row(s='b.c.d')]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.substring_index(_to_java_column(str),
delim, count))
-def levenshtein(left, right):
+def levenshtein(left: "ColumnOrName", right: "ColumnOrName") -> Column:
"""Computes the Levenshtein distance of the two given strings.
.. versionadded:: 1.5.0
@@ -2939,12 +3060,12 @@ def levenshtein(left, right):
>>> df0.select(levenshtein('l', 'r').alias('d')).collect()
[Row(d=3)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
jc = sc._jvm.functions.levenshtein(_to_java_column(left),
_to_java_column(right))
return Column(jc)
-def locate(substr, str, pos=1):
+def locate(substr: str, str: "ColumnOrName", pos: int = 1) -> Column:
"""
Locate the position of the first occurrence of substr in a string column,
after position pos.
@@ -2970,11 +3091,11 @@ def locate(substr, str, pos=1):
>>> df.select(locate('b', df.s, 1).alias('s')).collect()
[Row(s=2)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.locate(substr, _to_java_column(str), pos))
-def lpad(col, len, pad):
+def lpad(col: "ColumnOrName", len: int, pad: str) -> Column:
"""
Left-pad the string column to width `len` with `pad`.
@@ -2986,11 +3107,11 @@ def lpad(col, len, pad):
>>> df.select(lpad(df.s, 6, '#').alias('s')).collect()
[Row(s='##abcd')]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.lpad(_to_java_column(col), len, pad))
-def rpad(col, len, pad):
+def rpad(col: "ColumnOrName", len: int, pad: str) -> Column:
"""
Right-pad the string column to width `len` with `pad`.
@@ -3002,11 +3123,11 @@ def rpad(col, len, pad):
>>> df.select(rpad(df.s, 6, '#').alias('s')).collect()
[Row(s='abcd##')]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.rpad(_to_java_column(col), len, pad))
-def repeat(col, n):
+def repeat(col: "ColumnOrName", n: int) -> Column:
"""
Repeats a string column n times, and returns it as a new string column.
@@ -3018,11 +3139,11 @@ def repeat(col, n):
>>> df.select(repeat(df.s, 3).alias('s')).collect()
[Row(s='ababab')]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.repeat(_to_java_column(col), n))
-def split(str, pattern, limit=-1):
+def split(str: "ColumnOrName", pattern: str, limit: int = -1) -> Column:
"""
Splits str around matches of the given pattern.
@@ -3055,11 +3176,11 @@ def split(str, pattern, limit=-1):
>>> df.select(split(df.s, '[ABC]', -1).alias('s')).collect()
[Row(s=['one', 'two', 'three', ''])]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.split(_to_java_column(str), pattern,
limit))
-def regexp_extract(str, pattern, idx):
+def regexp_extract(str: "ColumnOrName", pattern: str, idx: int) -> Column:
r"""Extract a specific group matched by a Java regex, from the specified
string column.
If the regex did not match, or the specified group did not match, an empty
string is returned.
@@ -3077,12 +3198,12 @@ def regexp_extract(str, pattern, idx):
>>> df.select(regexp_extract('str', '(a+)(b)?(c)', 2).alias('d')).collect()
[Row(d='')]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
jc = sc._jvm.functions.regexp_extract(_to_java_column(str), pattern, idx)
return Column(jc)
-def regexp_replace(str, pattern, replacement):
+def regexp_replace(str: "ColumnOrName", pattern: str, replacement: str) ->
Column:
r"""Replace all substrings of the specified string value that match regexp
with rep.
.. versionadded:: 1.5.0
@@ -3093,12 +3214,12 @@ def regexp_replace(str, pattern, replacement):
>>> df.select(regexp_replace('str', r'(\d+)', '--').alias('d')).collect()
[Row(d='-----')]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
jc = sc._jvm.functions.regexp_replace(_to_java_column(str), pattern,
replacement)
return Column(jc)
-def initcap(col):
+def initcap(col: "ColumnOrName") -> Column:
"""Translate the first letter of each word to upper case in the sentence.
.. versionadded:: 1.5.0
@@ -3108,11 +3229,11 @@ def initcap(col):
>>> spark.createDataFrame([('ab cd',)],
['a']).select(initcap("a").alias('v')).collect()
[Row(v='Ab Cd')]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.initcap(_to_java_column(col)))
-def soundex(col):
+def soundex(col: "ColumnOrName") -> Column:
"""
Returns the SoundEx encoding for a string
@@ -3124,11 +3245,11 @@ def soundex(col):
>>> df.select(soundex(df.name).alias("soundex")).collect()
[Row(soundex='P362'), Row(soundex='U612')]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.soundex(_to_java_column(col)))
-def bin(col):
+def bin(col: "ColumnOrName") -> Column:
"""Returns the string representation of the binary value of the given
column.
.. versionadded:: 1.5.0
@@ -3138,12 +3259,12 @@ def bin(col):
>>> df.select(bin(df.age).alias('c')).collect()
[Row(c='10'), Row(c='101')]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
jc = sc._jvm.functions.bin(_to_java_column(col))
return Column(jc)
-def hex(col):
+def hex(col: "ColumnOrName") -> Column:
"""Computes hex value of the given column, which could be
:class:`pyspark.sql.types.StringType`,
:class:`pyspark.sql.types.BinaryType`,
:class:`pyspark.sql.types.IntegerType` or
:class:`pyspark.sql.types.LongType`.
@@ -3155,12 +3276,12 @@ def hex(col):
>>> spark.createDataFrame([('ABC', 3)], ['a', 'b']).select(hex('a'),
hex('b')).collect()
[Row(hex(a)='414243', hex(b)='3')]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
jc = sc._jvm.functions.hex(_to_java_column(col))
return Column(jc)
-def unhex(col):
+def unhex(col: "ColumnOrName") -> Column:
"""Inverse of hex. Interprets each pair of characters as a hexadecimal
number
and converts to the byte representation of number.
@@ -3171,11 +3292,11 @@ def unhex(col):
>>> spark.createDataFrame([('414243',)],
['a']).select(unhex('a')).collect()
[Row(unhex(a)=bytearray(b'ABC'))]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.unhex(_to_java_column(col)))
-def length(col):
+def length(col: "ColumnOrName") -> Column:
"""Computes the character length of string data or number of bytes of
binary data.
The length of character data includes the trailing spaces. The length of
binary data
includes binary zeros.
@@ -3187,11 +3308,11 @@ def length(col):
>>> spark.createDataFrame([('ABC ',)],
['a']).select(length('a').alias('length')).collect()
[Row(length=4)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.length(_to_java_column(col)))
-def octet_length(col):
+def octet_length(col: "ColumnOrName") -> Column:
"""
Calculates the byte length for the specified string column.
@@ -3217,7 +3338,7 @@ def octet_length(col):
return _invoke_function_over_column("octet_length", col)
-def bit_length(col):
+def bit_length(col: "ColumnOrName") -> Column:
"""
Calculates the bit length for the specified string column.
@@ -3243,7 +3364,7 @@ def bit_length(col):
return _invoke_function_over_column("bit_length", col)
-def translate(srcCol, matching, replace):
+def translate(srcCol: "ColumnOrName", matching: str, replace: str) -> Column:
"""A function translate any character in the `srcCol` by a character in
`matching`.
The characters in `replace` is corresponding to the characters in
`matching`.
The translate will happen when any character in the string matching with
the character
@@ -3257,13 +3378,13 @@ def translate(srcCol, matching, replace):
... .alias('r')).collect()
[Row(r='1a2s3ae')]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.translate(_to_java_column(srcCol),
matching, replace))
# ---------------------- Collection functions ------------------------------
-def create_map(*cols):
+def create_map(*cols: "ColumnOrName") -> Column:
"""Creates a new map column.
.. versionadded:: 2.0.0
@@ -3281,14 +3402,14 @@ def create_map(*cols):
>>> df.select(create_map([df.name, df.age]).alias("map")).collect()
[Row(map={'Alice': 2}), Row(map={'Bob': 5})]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
if len(cols) == 1 and isinstance(cols[0], (list, set)):
cols = cols[0]
jc = sc._jvm.functions.map(_to_seq(sc, cols, _to_java_column))
return Column(jc)
-def map_from_arrays(col1, col2):
+def map_from_arrays(col1: "ColumnOrName", col2: "ColumnOrName") -> Column:
"""Creates a new map from two arrays.
.. versionadded:: 2.4.0
@@ -3310,11 +3431,11 @@ def map_from_arrays(col1, col2):
|{2 -> a, 5 -> b}|
+----------------+
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.map_from_arrays(_to_java_column(col1),
_to_java_column(col2)))
-def array(*cols):
+def array(*cols: "ColumnOrName") -> Column:
"""Creates a new array column.
.. versionadded:: 1.4.0
@@ -3332,14 +3453,14 @@ def array(*cols):
>>> df.select(array([df.age, df.age]).alias("arr")).collect()
[Row(arr=[2, 2]), Row(arr=[5, 5])]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
if len(cols) == 1 and isinstance(cols[0], (list, set)):
cols = cols[0]
jc = sc._jvm.functions.array(_to_seq(sc, cols, _to_java_column))
return Column(jc)
-def array_contains(col, value):
+def array_contains(col: "ColumnOrName", value: Any) -> Column:
"""
Collection function: returns null if the array is null, true if the array
contains the
given value, and false otherwise.
@@ -3361,12 +3482,12 @@ def array_contains(col, value):
>>> df.select(array_contains(df.data, lit("a"))).collect()
[Row(array_contains(data, a)=True), Row(array_contains(data, a)=False)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
value = value._jc if isinstance(value, Column) else value
return Column(sc._jvm.functions.array_contains(_to_java_column(col),
value))
-def arrays_overlap(a1, a2):
+def arrays_overlap(a1: "ColumnOrName", a2: "ColumnOrName") -> Column:
"""
Collection function: returns true if the arrays contain any common
non-null element; if not,
returns null if both the arrays are non-empty and any of them contains a
null element; returns
@@ -3380,11 +3501,13 @@ def arrays_overlap(a1, a2):
>>> df.select(arrays_overlap(df.x, df.y).alias("overlap")).collect()
[Row(overlap=True), Row(overlap=False)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.arrays_overlap(_to_java_column(a1),
_to_java_column(a2)))
-def slice(x, start, length):
+def slice(
+ x: "ColumnOrName", start: Union[Column, int], length: Union[Column, int]
+) -> Column:
"""
Collection function: returns an array containing all the elements in `x`
from index `start`
(array indices start at 1, or from the end if `start` is negative) with
the specified `length`.
@@ -3406,7 +3529,7 @@ def slice(x, start, length):
>>> df.select(slice(df.x, 2, 2).alias("sliced")).collect()
[Row(sliced=[2, 3]), Row(sliced=[5])]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.slice(
_to_java_column(x),
start._jc if isinstance(start, Column) else start,
@@ -3414,7 +3537,9 @@ def slice(x, start, length):
))
-def array_join(col, delimiter, null_replacement=None):
+def array_join(
+ col: "ColumnOrName", delimiter: str, null_replacement: Optional[str] = None
+) -> Column:
"""
Concatenates the elements of `column` using the `delimiter`. Null values
are replaced with
`null_replacement` if set, otherwise they are ignored.
@@ -3429,7 +3554,7 @@ def array_join(col, delimiter, null_replacement=None):
>>> df.select(array_join(df.data, ",", "NULL").alias("joined")).collect()
[Row(joined='a,b,c'), Row(joined='a,NULL')]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
if null_replacement is None:
return Column(sc._jvm.functions.array_join(_to_java_column(col),
delimiter))
else:
@@ -3437,7 +3562,7 @@ def array_join(col, delimiter, null_replacement=None):
_to_java_column(col), delimiter, null_replacement))
-def concat(*cols):
+def concat(*cols: "ColumnOrName") -> Column:
"""
Concatenates multiple input columns together into a single column.
The function works with strings, binary and compatible array columns.
@@ -3454,11 +3579,11 @@ def concat(*cols):
>>> df.select(concat(df.a, df.b, df.c).alias("arr")).collect()
[Row(arr=[1, 2, 3, 4, 5]), Row(arr=None)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.concat(_to_seq(sc, cols, _to_java_column)))
-def array_position(col, value):
+def array_position(col: "ColumnOrName", value: Any) -> Column:
"""
Collection function: Locates the position of the first occurrence of the
given value
in the given array. Returns null if either of the arguments are null.
@@ -3476,11 +3601,11 @@ def array_position(col, value):
>>> df.select(array_position(df.data, "a")).collect()
[Row(array_position(data, a)=3), Row(array_position(data, a)=0)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.array_position(_to_java_column(col),
value))
-def element_at(col, extraction):
+def element_at(col: "ColumnOrName", extraction: Any) -> Column:
"""
Collection function: Returns element of array at given index in extraction
if col is array.
Returns value for the given key in extraction if col is map.
@@ -3508,12 +3633,12 @@ def element_at(col, extraction):
>>> df.select(element_at(df.data, lit("a"))).collect()
[Row(element_at(data, a)=1.0), Row(element_at(data, a)=None)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.element_at(
_to_java_column(col), lit(extraction)._jc))
-def array_remove(col, element):
+def array_remove(col: "ColumnOrName", element: Any) -> Column:
"""
Collection function: Remove all elements that equal to element from the
given array.
@@ -3532,11 +3657,11 @@ def array_remove(col, element):
>>> df.select(array_remove(df.data, 1)).collect()
[Row(array_remove(data, 1)=[2, 3]), Row(array_remove(data, 1)=[])]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.array_remove(_to_java_column(col),
element))
-def array_distinct(col):
+def array_distinct(col: "ColumnOrName") -> Column:
"""
Collection function: removes duplicate values from the array.
@@ -3553,11 +3678,11 @@ def array_distinct(col):
>>> df.select(array_distinct(df.data)).collect()
[Row(array_distinct(data)=[1, 2, 3]), Row(array_distinct(data)=[4, 5])]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.array_distinct(_to_java_column(col)))
-def array_intersect(col1, col2):
+def array_intersect(col1: "ColumnOrName", col2: "ColumnOrName") -> Column:
"""
Collection function: returns an array of the elements in the intersection
of col1 and col2,
without duplicates.
@@ -3578,11 +3703,11 @@ def array_intersect(col1, col2):
>>> df.select(array_intersect(df.c1, df.c2)).collect()
[Row(array_intersect(c1, c2)=['a', 'c'])]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.array_intersect(_to_java_column(col1),
_to_java_column(col2)))
-def array_union(col1, col2):
+def array_union(col1: "ColumnOrName", col2: "ColumnOrName") -> Column:
"""
Collection function: returns an array of the elements in the union of col1
and col2,
without duplicates.
@@ -3603,11 +3728,11 @@ def array_union(col1, col2):
>>> df.select(array_union(df.c1, df.c2)).collect()
[Row(array_union(c1, c2)=['b', 'a', 'c', 'd', 'f'])]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.array_union(_to_java_column(col1),
_to_java_column(col2)))
-def array_except(col1, col2):
+def array_except(col1: "ColumnOrName", col2: "ColumnOrName") -> Column:
"""
Collection function: returns an array of the elements in col1 but not in
col2,
without duplicates.
@@ -3628,11 +3753,11 @@ def array_except(col1, col2):
>>> df.select(array_except(df.c1, df.c2)).collect()
[Row(array_except(c1, c2)=['b'])]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.array_except(_to_java_column(col1),
_to_java_column(col2)))
-def explode(col):
+def explode(col: "ColumnOrName") -> Column:
"""
Returns a new row for each element in the given array or map.
Uses the default column name `col` for elements in the array and
@@ -3654,12 +3779,12 @@ def explode(col):
| a| b|
+---+-----+
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
jc = sc._jvm.functions.explode(_to_java_column(col))
return Column(jc)
-def posexplode(col):
+def posexplode(col: "ColumnOrName") -> Column:
"""
Returns a new row for each element with position in the given array or map.
Uses the default column name `pos` for position, and `col` for elements in
the
@@ -3681,12 +3806,12 @@ def posexplode(col):
| 0| a| b|
+---+---+-----+
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
jc = sc._jvm.functions.posexplode(_to_java_column(col))
return Column(jc)
-def explode_outer(col):
+def explode_outer(col: "ColumnOrName") -> Column:
"""
Returns a new row for each element in the given array or map.
Unlike explode, if the array/map is null or empty then null is produced.
@@ -3720,12 +3845,12 @@ def explode_outer(col):
| 3| null|null|
+---+----------+----+
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
jc = sc._jvm.functions.explode_outer(_to_java_column(col))
return Column(jc)
-def posexplode_outer(col):
+def posexplode_outer(col: "ColumnOrName") -> Column:
"""
Returns a new row for each element with position in the given array or map.
Unlike posexplode, if the array/map is null or empty then the row (null,
null) is produced.
@@ -3758,12 +3883,12 @@ def posexplode_outer(col):
| 3| null|null|null|
+---+----------+----+----+
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
jc = sc._jvm.functions.posexplode_outer(_to_java_column(col))
return Column(jc)
-def get_json_object(col, path):
+def get_json_object(col: "ColumnOrName", path: str) -> Column:
"""
Extracts json object from a json string based on json path specified, and
returns json string
of the extracted json object. It will return null if the input json string
is invalid.
@@ -3785,12 +3910,12 @@ def get_json_object(col, path):
... get_json_object(df.jstring, '$.f2').alias("c1")
).collect()
[Row(key='1', c0='value1', c1='value2'), Row(key='2', c0='value12',
c1=None)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
jc = sc._jvm.functions.get_json_object(_to_java_column(col), path)
return Column(jc)
-def json_tuple(col, *fields):
+def json_tuple(col: "ColumnOrName", *fields: str) -> Column:
"""Creates a new row for a json column according to the given field names.
.. versionadded:: 1.6.0
@@ -3809,12 +3934,16 @@ def json_tuple(col, *fields):
>>> df.select(df.key, json_tuple(df.jstring, 'f1', 'f2')).collect()
[Row(key='1', c0='value1', c1='value2'), Row(key='2', c0='value12',
c1=None)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
jc = sc._jvm.functions.json_tuple(_to_java_column(col), _to_seq(sc,
fields))
return Column(jc)
-def from_json(col, schema, options=None):
+def from_json(
+ col: "ColumnOrName",
+ schema: Union[ArrayType, StructType, Column, str],
+ options: Optional[Dict[str, str]] = None,
+) -> Column:
"""
Parses a column containing a JSON string into a :class:`MapType` with
:class:`StringType`
as keys type, :class:`StructType` or :class:`ArrayType` with
@@ -3865,7 +3994,7 @@ def from_json(col, schema, options=None):
[Row(json=[1, 2, 3])]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
if isinstance(schema, DataType):
schema = schema.json()
elif isinstance(schema, Column):
@@ -3874,7 +4003,7 @@ def from_json(col, schema, options=None):
return Column(jc)
-def to_json(col, options=None):
+def to_json(col: "ColumnOrName", options: Optional[Dict[str, str]] = None) ->
Column:
"""
Converts a column containing a :class:`StructType`, :class:`ArrayType` or
a :class:`MapType`
into a JSON string. Throws an exception, in the case of an unsupported
type.
@@ -3920,12 +4049,12 @@ def to_json(col, options=None):
[Row(json='["Alice","Bob"]')]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
jc = sc._jvm.functions.to_json(_to_java_column(col),
_options_to_str(options))
return Column(jc)
-def schema_of_json(json, options=None):
+def schema_of_json(json: "ColumnOrName", options: Optional[Dict[str, str]] =
None) -> Column:
"""
Parses a JSON string and infers its schema in DDL format.
@@ -3961,12 +4090,12 @@ def schema_of_json(json, options=None):
else:
raise TypeError("schema argument should be a column or string")
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
jc = sc._jvm.functions.schema_of_json(col, _options_to_str(options))
return Column(jc)
-def schema_of_csv(csv, options=None):
+def schema_of_csv(csv: "ColumnOrName", options: Optional[Dict[str, str]] =
None) -> Column:
"""
Parses a CSV string and infers its schema in DDL format.
@@ -3998,12 +4127,12 @@ def schema_of_csv(csv, options=None):
else:
raise TypeError("schema argument should be a column or string")
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
jc = sc._jvm.functions.schema_of_csv(col, _options_to_str(options))
return Column(jc)
-def to_csv(col, options=None):
+def to_csv(col: "ColumnOrName", options: Optional[Dict[str, str]] = None) ->
Column:
"""
Converts a column containing a :class:`StructType` into a CSV string.
Throws an exception, in the case of an unsupported type.
@@ -4030,12 +4159,12 @@ def to_csv(col, options=None):
[Row(csv='2,Alice')]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
jc = sc._jvm.functions.to_csv(_to_java_column(col),
_options_to_str(options))
return Column(jc)
-def size(col):
+def size(col: "ColumnOrName") -> Column:
"""
Collection function: returns the length of the array or map stored in the
column.
@@ -4052,11 +4181,11 @@ def size(col):
>>> df.select(size(df.data)).collect()
[Row(size(data)=3), Row(size(data)=1), Row(size(data)=0)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.size(_to_java_column(col)))
-def array_min(col):
+def array_min(col: "ColumnOrName") -> Column:
"""
Collection function: returns the minimum value of the array.
@@ -4073,11 +4202,11 @@ def array_min(col):
>>> df.select(array_min(df.data).alias('min')).collect()
[Row(min=1), Row(min=-1)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.array_min(_to_java_column(col)))
-def array_max(col):
+def array_max(col: "ColumnOrName") -> Column:
"""
Collection function: returns the maximum value of the array.
@@ -4094,11 +4223,11 @@ def array_max(col):
>>> df.select(array_max(df.data).alias('max')).collect()
[Row(max=3), Row(max=10)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.array_max(_to_java_column(col)))
-def sort_array(col, asc=True):
+def sort_array(col: "ColumnOrName", asc: bool = True) -> Column:
"""
Collection function: sorts the input array in ascending or descending
order according
to the natural ordering of the array elements. Null elements will be
placed at the beginning
@@ -4121,11 +4250,11 @@ def sort_array(col, asc=True):
>>> df.select(sort_array(df.data, asc=False).alias('r')).collect()
[Row(r=[3, 2, 1, None]), Row(r=[1]), Row(r=[])]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.sort_array(_to_java_column(col), asc))
-def array_sort(col):
+def array_sort(col: "ColumnOrName") -> Column:
"""
Collection function: sorts the input array in ascending order. The
elements of the input array
must be orderable. Null elements will be placed at the end of the returned
array.
@@ -4143,11 +4272,11 @@ def array_sort(col):
>>> df.select(array_sort(df.data).alias('r')).collect()
[Row(r=[1, 2, 3, None]), Row(r=[1]), Row(r=[])]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.array_sort(_to_java_column(col)))
-def shuffle(col):
+def shuffle(col: "ColumnOrName") -> Column:
"""
Collection function: Generates a random permutation of the given array.
@@ -4168,11 +4297,11 @@ def shuffle(col):
>>> df.select(shuffle(df.data).alias('s')).collect() # doctest: +SKIP
[Row(s=[3, 1, 5, 20]), Row(s=[20, None, 3, 1])]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.shuffle(_to_java_column(col)))
-def reverse(col):
+def reverse(col: "ColumnOrName") -> Column:
"""
Collection function: returns a reversed string or an array with reverse
order of elements.
@@ -4192,11 +4321,11 @@ def reverse(col):
>>> df.select(reverse(df.data).alias('r')).collect()
[Row(r=[3, 1, 2]), Row(r=[1]), Row(r=[])]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.reverse(_to_java_column(col)))
-def flatten(col):
+def flatten(col: "ColumnOrName") -> Column:
"""
Collection function: creates a single array from an array of arrays.
If a structure of nested arrays is deeper than two levels,
@@ -4215,11 +4344,11 @@ def flatten(col):
>>> df.select(flatten(df.data).alias('r')).collect()
[Row(r=[1, 2, 3, 4, 5, 6]), Row(r=None)]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.flatten(_to_java_column(col)))
-def map_keys(col):
+def map_keys(col: "ColumnOrName") -> Column:
"""
Collection function: Returns an unordered array containing the keys of the
map.
@@ -4241,11 +4370,11 @@ def map_keys(col):
|[1, 2]|
+------+
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.map_keys(_to_java_column(col)))
-def map_values(col):
+def map_values(col: "ColumnOrName") -> Column:
"""
Collection function: Returns an unordered array containing the values of
the map.
@@ -4267,11 +4396,11 @@ def map_values(col):
|[a, b]|
+------+
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.map_values(_to_java_column(col)))
-def map_entries(col):
+def map_entries(col: "ColumnOrName") -> Column:
"""
Collection function: Returns an unordered array of all entries in the
given map.
@@ -4293,11 +4422,11 @@ def map_entries(col):
|[{1, a}, {2, b}]|
+----------------+
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.map_entries(_to_java_column(col)))
-def map_from_entries(col):
+def map_from_entries(col: "ColumnOrName") -> Column:
"""
Collection function: Returns a map created from the given array of entries.
@@ -4319,11 +4448,11 @@ def map_from_entries(col):
|{1 -> a, 2 -> b}|
+----------------+
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.map_from_entries(_to_java_column(col)))
-def array_repeat(col, count):
+def array_repeat(col: "ColumnOrName", count: Union[Column, int]) -> Column:
"""
Collection function: creates an array containing a column repeated count
times.
@@ -4335,14 +4464,14 @@ def array_repeat(col, count):
>>> df.select(array_repeat(df.data, 3).alias('r')).collect()
[Row(r=['ab', 'ab', 'ab'])]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.array_repeat(
_to_java_column(col),
_to_java_column(count) if isinstance(count, Column) else count
))
-def arrays_zip(*cols):
+def arrays_zip(*cols: "ColumnOrName") -> Column:
"""
Collection function: Returns a merged array of structs in which the N-th
struct contains all
N-th values of input arrays.
@@ -4361,11 +4490,11 @@ def arrays_zip(*cols):
>>> df.select(arrays_zip(df.vals1, df.vals2).alias('zipped')).collect()
[Row(zipped=[Row(vals1=1, vals2=2), Row(vals1=2, vals2=3), Row(vals1=3,
vals2=4)])]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.arrays_zip(_to_seq(sc, cols,
_to_java_column)))
-def map_concat(*cols):
+def map_concat(*cols: "ColumnOrName") -> Column:
"""Returns the union of all the given maps.
.. versionadded:: 2.4.0
@@ -4386,14 +4515,16 @@ def map_concat(*cols):
|{1 -> a, 2 -> b, 3 -> c}|
+------------------------+
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
if len(cols) == 1 and isinstance(cols[0], (list, set)):
cols = cols[0]
jc = sc._jvm.functions.map_concat(_to_seq(sc, cols, _to_java_column))
return Column(jc)
-def sequence(start, stop, step=None):
+def sequence(
+ start: "ColumnOrName", stop: "ColumnOrName", step:
Optional["ColumnOrName"] = None
+) -> Column:
"""
Generate a sequence of integers from `start` to `stop`, incrementing by
`step`.
If `step` is not set, incrementing by 1 if `start` is less than or equal
to `stop`,
@@ -4410,7 +4541,7 @@ def sequence(start, stop, step=None):
>>> df2.select(sequence('C1', 'C2', 'C3').alias('r')).collect()
[Row(r=[4, 2, 0, -2, -4])]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
if step is None:
return Column(sc._jvm.functions.sequence(_to_java_column(start),
_to_java_column(stop)))
else:
@@ -4418,7 +4549,11 @@ def sequence(start, stop, step=None):
_to_java_column(start), _to_java_column(stop),
_to_java_column(step)))
-def from_csv(col, schema, options=None):
+def from_csv(
+ col: "ColumnOrName",
+ schema: Union[StructType, Column, str],
+ options: Optional[Dict[str, str]] = None,
+) -> Column:
"""
Parses a column containing a CSV string to a row with the specified schema.
Returns `null`, in the case of an unparseable string.
@@ -4454,7 +4589,7 @@ def from_csv(col, schema, options=None):
[Row(csv=Row(s='abc'))]
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
if isinstance(schema, str):
schema = _create_column_from_literal(schema)
elif isinstance(schema, Column):
@@ -4466,7 +4601,7 @@ def from_csv(col, schema, options=None):
return Column(jc)
-def _unresolved_named_lambda_variable(*name_parts):
+def _unresolved_named_lambda_variable(*name_parts: Any) -> Column:
"""
Create `o.a.s.sql.expressions.UnresolvedNamedLambdaVariable`,
convert it to o.s.sql.Column and wrap in Python `Column`
@@ -4475,7 +4610,7 @@ def _unresolved_named_lambda_variable(*name_parts):
----------
name_parts : str
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
name_parts_seq = _to_seq(sc, name_parts)
expressions = sc._jvm.org.apache.spark.sql.catalyst.expressions
return Column(
@@ -4485,9 +4620,7 @@ def _unresolved_named_lambda_variable(*name_parts):
)
-def _get_lambda_parameters(f):
- import inspect
-
+def _get_lambda_parameters(f: Callable) -> ValuesView[inspect.Parameter]:
signature = inspect.signature(f)
parameters = signature.parameters.values()
@@ -4517,7 +4650,7 @@ def _get_lambda_parameters(f):
return parameters
-def _create_lambda(f):
+def _create_lambda(f: Callable) -> Callable:
"""
Create `o.a.s.sql.expressions.LambdaFunction` corresponding
to transformation described by f
@@ -4529,7 +4662,7 @@ def _create_lambda(f):
"""
parameters = _get_lambda_parameters(f)
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
expressions = sc._jvm.org.apache.spark.sql.catalyst.expressions
argnames = ["x", "y", "z"]
@@ -4545,13 +4678,17 @@ def _create_lambda(f):
if not isinstance(result, Column):
raise ValueError("f should return Column, got {}".format(type(result)))
- jexpr = result._jc.expr()
- jargs = _to_seq(sc, [arg._jc.expr() for arg in args])
+ jexpr = result._jc.expr() # type: ignore[operator]
+ jargs = _to_seq(sc, [arg._jc.expr() for arg in args]) # type:
ignore[operator]
return expressions.LambdaFunction(jexpr, jargs, False)
-def _invoke_higher_order_function(name, cols, funs):
+def _invoke_higher_order_function(
+ name: str,
+ cols: List["ColumnOrName"],
+ funs: List[Callable],
+) -> Column:
"""
Invokes expression identified by name,
(relative to ```org.apache.spark.sql.catalyst.expressions``)
@@ -4563,7 +4700,7 @@ def _invoke_higher_order_function(name, cols, funs):
:return: a Column
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
expressions = sc._jvm.org.apache.spark.sql.catalyst.expressions
expr = getattr(expressions, name)
@@ -4573,7 +4710,20 @@ def _invoke_higher_order_function(name, cols, funs):
return Column(sc._jvm.Column(expr(*jcols + jfuns)))
-def transform(col, f):
+@overload
+def transform(col: "ColumnOrName", f: Callable[[Column], Column]) -> Column:
+ ...
+
+
+@overload
+def transform(col: "ColumnOrName", f: Callable[[Column, Column], Column]) ->
Column:
+ ...
+
+
+def transform(
+ col: "ColumnOrName",
+ f: Union[Callable[[Column], Column], Callable[[Column, Column], Column]],
+) -> Column:
"""
Returns an array of elements after applying a transformation to each
element in the input array.
@@ -4622,7 +4772,7 @@ def transform(col, f):
return _invoke_higher_order_function("ArrayTransform", [col], [f])
-def exists(col, f):
+def exists(col: "ColumnOrName", f: Callable[[Column], Column]) -> Column:
"""
Returns whether a predicate holds for one or more elements in the array.
@@ -4654,7 +4804,7 @@ def exists(col, f):
return _invoke_higher_order_function("ArrayExists", [col], [f])
-def forall(col, f):
+def forall(col: "ColumnOrName", f: Callable[[Column], Column]) -> Column:
"""
Returns whether a predicate holds for every element in the array.
@@ -4693,7 +4843,20 @@ def forall(col, f):
return _invoke_higher_order_function("ArrayForAll", [col], [f])
-def filter(col, f):
+@overload
+def filter(col: "ColumnOrName", f: Callable[[Column], Column]) -> Column:
+ ...
+
+
+@overload
+def filter(col: "ColumnOrName", f: Callable[[Column, Column], Column]) ->
Column:
+ ...
+
+
+def filter(
+ col: "ColumnOrName",
+ f: Union[Callable[[Column], Column], Callable[[Column, Column], Column]],
+) -> Column:
"""
Returns an array of elements for which a predicate holds in a given array.
@@ -4740,7 +4903,12 @@ def filter(col, f):
return _invoke_higher_order_function("ArrayFilter", [col], [f])
-def aggregate(col, initialValue, merge, finish=None):
+def aggregate(
+ col: "ColumnOrName",
+ initialValue: "ColumnOrName",
+ merge: Callable[[Column, Column], Column],
+ finish: Optional[Callable[[Column], Column]] = None,
+) -> Column:
"""
Applies a binary operator to an initial state and all elements in the
array,
and reduces this to a single state. The final state is converted into the
final result
@@ -4813,7 +4981,11 @@ def aggregate(col, initialValue, merge, finish=None):
)
-def zip_with(left, right, f):
+def zip_with(
+ left: "ColumnOrName",
+ right: "ColumnOrName",
+ f: Callable[[Column, Column], Column],
+) -> Column:
"""
Merge two given arrays, element-wise, into a single array using a function.
If one array is shorter, nulls are appended at the end to match the length
of the longer
@@ -4859,7 +5031,9 @@ def zip_with(left, right, f):
return _invoke_higher_order_function("ZipWith", [left, right], [f])
-def transform_keys(col, f):
+def transform_keys(
+ col: "ColumnOrName", f: Callable[[Column, Column], Column]
+) -> Column:
"""
Applies a function to every key-value pair in a map and returns
a map with the results of those applications as the new keys for the pairs.
@@ -4896,7 +5070,9 @@ def transform_keys(col, f):
return _invoke_higher_order_function("TransformKeys", [col], [f])
-def transform_values(col, f):
+def transform_values(
+ col: "ColumnOrName", f: Callable[[Column, Column], Column]
+) -> Column:
"""
Applies a function to every key-value pair in a map and returns
a map with the results of those applications as the new values for the
pairs.
@@ -4933,7 +5109,7 @@ def transform_values(col, f):
return _invoke_higher_order_function("TransformValues", [col], [f])
-def map_filter(col, f):
+def map_filter(col: "ColumnOrName", f: Callable[[Column, Column], Column]) ->
Column:
"""
Returns a map whose key-value pairs satisfy a predicate.
@@ -4969,7 +5145,11 @@ def map_filter(col, f):
return _invoke_higher_order_function("MapFilter", [col], [f])
-def map_zip_with(col1, col2, f):
+def map_zip_with(
+ col1: "ColumnOrName",
+ col2: "ColumnOrName",
+ f: Callable[[Column, Column, Column], Column],
+) -> Column:
"""
Merge two given maps, key-wise into a single map using a function.
@@ -5012,7 +5192,7 @@ def map_zip_with(col1, col2, f):
# ---------------------- Partition transform functions
--------------------------------
-def years(col):
+def years(col: "ColumnOrName") -> Column:
"""
Partition transform function: A transform for timestamps and dates
to partition data into years.
@@ -5032,11 +5212,11 @@ def years(col):
method of the `DataFrameWriterV2`.
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.years(_to_java_column(col)))
-def months(col):
+def months(col: "ColumnOrName") -> Column:
"""
Partition transform function: A transform for timestamps and dates
to partition data into months.
@@ -5056,11 +5236,11 @@ def months(col):
method of the `DataFrameWriterV2`.
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.months(_to_java_column(col)))
-def days(col):
+def days(col: "ColumnOrName") -> Column:
"""
Partition transform function: A transform for timestamps and dates
to partition data into days.
@@ -5080,11 +5260,11 @@ def days(col):
method of the `DataFrameWriterV2`.
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.days(_to_java_column(col)))
-def hours(col):
+def hours(col: "ColumnOrName") -> Column:
"""
Partition transform function: A transform for timestamps
to partition data into hours.
@@ -5104,11 +5284,11 @@ def hours(col):
method of the `DataFrameWriterV2`.
"""
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
return Column(sc._jvm.functions.hours(_to_java_column(col)))
-def bucket(numBuckets, col):
+def bucket(numBuckets: Union[Column, int], col: "ColumnOrName") -> Column:
"""
Partition transform function: A transform for any type that partitions
by a hash of the input column.
@@ -5133,7 +5313,7 @@ def bucket(numBuckets, col):
"numBuckets should be a Column or an int, got
{}".format(type(numBuckets))
)
- sc = SparkContext._active_spark_context
+ sc = SparkContext._active_spark_context # type: ignore[attr-defined]
numBuckets = (
_create_column_from_literal(numBuckets)
if isinstance(numBuckets, int)
@@ -5144,7 +5324,32 @@ def bucket(numBuckets, col):
# ---------------------------- User Defined Function
----------------------------------
-def udf(f=None, returnType=StringType()):
+@overload
+def udf(
+ f: Callable[..., Any], returnType: "DataTypeOrString" = StringType()
+) -> "UserDefinedFunctionLike":
+ ...
+
+
+@overload
+def udf(
+ f: Optional["DataTypeOrString"] = None,
+) -> Callable[[Callable[..., Any]], "UserDefinedFunctionLike"]:
+ ...
+
+
+@overload
+def udf(
+ *,
+ returnType: "DataTypeOrString" = StringType(),
+) -> Callable[[Callable[..., Any]], "UserDefinedFunctionLike"]:
+ ...
+
+
+def udf(
+ f: Optional[Union[Callable[..., Any], "DataTypeOrString"]] = None,
+ returnType: Optional["DataTypeOrString"] = StringType(),
+) -> Union["UserDefinedFunctionLike", Callable[[Callable[..., Any]],
"UserDefinedFunctionLike"]]:
"""Creates a user defined function (UDF).
.. versionadded:: 1.3.0
@@ -5240,7 +5445,7 @@ def udf(f=None, returnType=StringType()):
evalType=PythonEvalType.SQL_BATCHED_UDF)
-def _test():
+def _test() -> None:
import doctest
from pyspark.sql import Row, SparkSession
import pyspark.sql.functions
diff --git a/python/pyspark/sql/functions.pyi b/python/pyspark/sql/functions.pyi
deleted file mode 100644
index a147210..0000000
--- a/python/pyspark/sql/functions.pyi
+++ /dev/null
@@ -1,377 +0,0 @@
-#
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-from typing import overload
-from typing import Any, Callable, Dict, List, Optional, Union
-
-from pyspark.sql._typing import (
- ColumnOrName,
- DataTypeOrString,
- UserDefinedFunctionLike,
-)
-from pyspark.sql.pandas.functions import ( # noqa: F401
- pandas_udf as pandas_udf,
- PandasUDFType as PandasUDFType,
-)
-from pyspark.sql.column import Column
-from pyspark.sql.dataframe import DataFrame
-from pyspark.sql.types import ( # noqa: F401
- ArrayType,
- StringType,
- StructType,
- DataType,
-)
-from pyspark.sql.utils import to_str # noqa: F401
-
-def approxCountDistinct(col: ColumnOrName, rsd: Optional[float] = ...) ->
Column: ...
-def approx_count_distinct(col: ColumnOrName, rsd: Optional[float] = ...) ->
Column: ...
-def broadcast(df: DataFrame) -> DataFrame: ...
-def coalesce(*cols: ColumnOrName) -> Column: ...
-def corr(col1: ColumnOrName, col2: ColumnOrName) -> Column: ...
-def covar_pop(col1: ColumnOrName, col2: ColumnOrName) -> Column: ...
-def covar_samp(col1: ColumnOrName, col2: ColumnOrName) -> Column: ...
-def countDistinct(col: ColumnOrName, *cols: ColumnOrName) -> Column: ...
-def count_distinct(col: ColumnOrName, *cols: ColumnOrName) -> Column: ...
-def first(col: ColumnOrName, ignorenulls: bool = ...) -> Column: ...
-def grouping(col: ColumnOrName) -> Column: ...
-def grouping_id(*cols: ColumnOrName) -> Column: ...
-def input_file_name() -> Column: ...
-def isnan(col: ColumnOrName) -> Column: ...
-def isnull(col: ColumnOrName) -> Column: ...
-def last(col: ColumnOrName, ignorenulls: bool = ...) -> Column: ...
-def monotonically_increasing_id() -> Column: ...
-def nanvl(col1: ColumnOrName, col2: ColumnOrName) -> Column: ...
-def percentile_approx(
- col: ColumnOrName,
- percentage: Union[Column, float, List[float], tuple[float]],
- accuracy: Union[Column, float] = ...,
-) -> Column: ...
-def rand(seed: Optional[int] = ...) -> Column: ...
-def randn(seed: Optional[int] = ...) -> Column: ...
-def round(col: ColumnOrName, scale: int = ...) -> Column: ...
-def bround(col: ColumnOrName, scale: int = ...) -> Column: ...
-def shiftLeft(col: ColumnOrName, numBits: int) -> Column: ...
-def shiftleft(col: ColumnOrName, numBits: int) -> Column: ...
-def shiftRight(col: ColumnOrName, numBits: int) -> Column: ...
-def shiftright(col: ColumnOrName, numBits: int) -> Column: ...
-def shiftRightUnsigned(col: ColumnOrName, numBits: int) -> Column: ...
-def shiftrightunsigned(col: ColumnOrName, numBits: int) -> Column: ...
-def spark_partition_id() -> Column: ...
-def expr(str: str) -> Column: ...
-def struct(*cols: ColumnOrName) -> Column: ...
-def greatest(*cols: ColumnOrName) -> Column: ...
-def least(*cols: Column) -> Column: ...
-def when(condition: Column, value: Any) -> Column: ...
-@overload
-def log(arg1: ColumnOrName) -> Column: ...
-@overload
-def log(arg1: float, arg2: ColumnOrName) -> Column: ...
-def log2(col: ColumnOrName) -> Column: ...
-def conv(col: ColumnOrName, fromBase: int, toBase: int) -> Column: ...
-def factorial(col: ColumnOrName) -> Column: ...
-def lag(
- col: ColumnOrName, offset: int = ..., default: Optional[Any] = ...
-) -> Column: ...
-def lead(
- col: ColumnOrName, offset: int = ..., default: Optional[Any] = ...
-) -> Column: ...
-def nth_value(
- col: ColumnOrName, offset: int, ignoreNulls: Optional[bool] = ...
-) -> Column: ...
-def ntile(n: int) -> Column: ...
-def current_date() -> Column: ...
-def current_timestamp() -> Column: ...
-def date_format(date: ColumnOrName, format: str) -> Column: ...
-def year(col: ColumnOrName) -> Column: ...
-def quarter(col: ColumnOrName) -> Column: ...
-def month(col: ColumnOrName) -> Column: ...
-def dayofweek(col: ColumnOrName) -> Column: ...
-def dayofmonth(col: ColumnOrName) -> Column: ...
-def dayofyear(col: ColumnOrName) -> Column: ...
-def hour(col: ColumnOrName) -> Column: ...
-def minute(col: ColumnOrName) -> Column: ...
-def second(col: ColumnOrName) -> Column: ...
-def weekofyear(col: ColumnOrName) -> Column: ...
-def date_add(start: ColumnOrName, days: int) -> Column: ...
-def date_sub(start: ColumnOrName, days: int) -> Column: ...
-def datediff(end: ColumnOrName, start: ColumnOrName) -> Column: ...
-def add_months(start: ColumnOrName, months: int) -> Column: ...
-def months_between(
- date1: ColumnOrName, date2: ColumnOrName, roundOff: bool = ...
-) -> Column: ...
-def to_date(col: ColumnOrName, format: Optional[str] = ...) -> Column: ...
-@overload
-def to_timestamp(col: ColumnOrName) -> Column: ...
-@overload
-def to_timestamp(col: ColumnOrName, format: str) -> Column: ...
-def trunc(date: ColumnOrName, format: str) -> Column: ...
-def date_trunc(format: str, timestamp: ColumnOrName) -> Column: ...
-def next_day(date: ColumnOrName, dayOfWeek: str) -> Column: ...
-def last_day(date: ColumnOrName) -> Column: ...
-def from_unixtime(timestamp: ColumnOrName, format: str = ...) -> Column: ...
-def unix_timestamp(
- timestamp: Optional[ColumnOrName] = ..., format: str = ...
-) -> Column: ...
-def from_utc_timestamp(timestamp: ColumnOrName, tz: ColumnOrName) -> Column:
...
-def to_utc_timestamp(timestamp: ColumnOrName, tz: ColumnOrName) -> Column: ...
-def timestamp_seconds(col: ColumnOrName) -> Column: ...
-def window(
- timeColumn: ColumnOrName,
- windowDuration: str,
- slideDuration: Optional[str] = ...,
- startTime: Optional[str] = ...,
-) -> Column: ...
-def session_window(timeColumn: ColumnOrName, gapDuration: Union[Column, str])
-> Column: ...
-def crc32(col: ColumnOrName) -> Column: ...
-def md5(col: ColumnOrName) -> Column: ...
-def sha1(col: ColumnOrName) -> Column: ...
-def sha2(col: ColumnOrName, numBits: int) -> Column: ...
-def hash(*cols: ColumnOrName) -> Column: ...
-def xxhash64(*cols: ColumnOrName) -> Column: ...
-def assert_true(col: ColumnOrName, errMsg: Union[Column, str] = ...) ->
Column: ...
-def raise_error(errMsg: Union[Column, str]) -> Column: ...
-def concat(*cols: ColumnOrName) -> Column: ...
-def concat_ws(sep: str, *cols: ColumnOrName) -> Column: ...
-def decode(col: ColumnOrName, charset: str) -> Column: ...
-def encode(col: ColumnOrName, charset: str) -> Column: ...
-def format_number(col: ColumnOrName, d: int) -> Column: ...
-def format_string(format: str, *cols: ColumnOrName) -> Column: ...
-def instr(str: ColumnOrName, substr: str) -> Column: ...
-def overlay(
- src: ColumnOrName,
- replace: ColumnOrName,
- pos: Union[Column, int],
- len: Union[Column, int] = ...,
-) -> Column: ...
-def substring(str: ColumnOrName, pos: int, len: int) -> Column: ...
-def substring_index(str: ColumnOrName, delim: str, count: int) -> Column: ...
-def levenshtein(left: ColumnOrName, right: ColumnOrName) -> Column: ...
-def locate(substr: str, str: ColumnOrName, pos: int = ...) -> Column: ...
-def lpad(col: ColumnOrName, len: int, pad: str) -> Column: ...
-def rpad(col: ColumnOrName, len: int, pad: str) -> Column: ...
-def repeat(col: ColumnOrName, n: int) -> Column: ...
-def split(str: ColumnOrName, pattern: str, limit: int = ...) -> Column: ...
-def regexp_extract(str: ColumnOrName, pattern: str, idx: int) -> Column: ...
-def regexp_replace(str: ColumnOrName, pattern: str, replacement: str) ->
Column: ...
-def initcap(col: ColumnOrName) -> Column: ...
-def soundex(col: ColumnOrName) -> Column: ...
-def bin(col: ColumnOrName) -> Column: ...
-def hex(col: ColumnOrName) -> Column: ...
-def unhex(col: ColumnOrName) -> Column: ...
-def length(col: ColumnOrName) -> Column: ...
-def octet_length(col: ColumnOrName) -> Column: ...
-def bit_length(col: ColumnOrName) -> Column: ...
-def translate(srcCol: ColumnOrName, matching: str, replace: str) -> Column: ...
-def map_from_arrays(col1: ColumnOrName, col2: ColumnOrName) -> Column: ...
-def create_map(*cols: ColumnOrName) -> Column: ...
-def array(*cols: ColumnOrName) -> Column: ...
-def array_contains(col: ColumnOrName, value: Any) -> Column: ...
-def arrays_overlap(a1: ColumnOrName, a2: ColumnOrName) -> Column: ...
-def slice(
- x: ColumnOrName, start: Union[Column, int], length: Union[Column, int]
-) -> Column: ...
-def array_join(
- col: ColumnOrName, delimiter: str, null_replacement: Optional[str] = ...
-) -> Column: ...
-def array_position(col: ColumnOrName, value: Any) -> Column: ...
-def element_at(col: ColumnOrName, extraction: Any) -> Column: ...
-def array_remove(col: ColumnOrName, element: Any) -> Column: ...
-def array_distinct(col: ColumnOrName) -> Column: ...
-def array_intersect(col1: ColumnOrName, col2: ColumnOrName) -> Column: ...
-def array_union(col1: ColumnOrName, col2: ColumnOrName) -> Column: ...
-def array_except(col1: ColumnOrName, col2: ColumnOrName) -> Column: ...
-def explode(col: ColumnOrName) -> Column: ...
-def explode_outer(col: ColumnOrName) -> Column: ...
-def posexplode(col: ColumnOrName) -> Column: ...
-def posexplode_outer(col: ColumnOrName) -> Column: ...
-def get_json_object(col: ColumnOrName, path: str) -> Column: ...
-def json_tuple(col: ColumnOrName, *fields: str) -> Column: ...
-def from_json(
- col: ColumnOrName,
- schema: Union[ArrayType, StructType, Column, str],
- options: Optional[Dict[str, str]] = ...,
-) -> Column: ...
-def to_json(col: ColumnOrName, options: Optional[Dict[str, str]] = ...) ->
Column: ...
-def schema_of_json(json: ColumnOrName, options: Optional[Dict[str, str]] =
...) -> Column: ...
-def schema_of_csv(csv: ColumnOrName, options: Optional[Dict[str, str]] = ...)
-> Column: ...
-def to_csv(col: ColumnOrName, options: Optional[Dict[str, str]] = ...) ->
Column: ...
-def size(col: ColumnOrName) -> Column: ...
-def array_min(col: ColumnOrName) -> Column: ...
-def array_max(col: ColumnOrName) -> Column: ...
-def sort_array(col: ColumnOrName, asc: bool = ...) -> Column: ...
-def array_sort(col: ColumnOrName) -> Column: ...
-def shuffle(col: ColumnOrName) -> Column: ...
-def reverse(col: ColumnOrName) -> Column: ...
-def flatten(col: ColumnOrName) -> Column: ...
-def map_keys(col: ColumnOrName) -> Column: ...
-def map_values(col: ColumnOrName) -> Column: ...
-def map_entries(col: ColumnOrName) -> Column: ...
-def map_from_entries(col: ColumnOrName) -> Column: ...
-def array_repeat(col: ColumnOrName, count: Union[Column, int]) -> Column: ...
-def arrays_zip(*cols: ColumnOrName) -> Column: ...
-def map_concat(*cols: ColumnOrName) -> Column: ...
-def sentences(
- col: ColumnOrName,
- language: Optional[ColumnOrName] = ...,
- country: Optional[ColumnOrName] = ...
-) -> Column: ...
-def sequence(
- start: ColumnOrName, stop: ColumnOrName, step: Optional[ColumnOrName] = ...
-) -> Column: ...
-def from_csv(
- col: ColumnOrName,
- schema: Union[StructType, Column, str],
- options: Optional[Dict[str, str]] = ...,
-) -> Column: ...
-@overload
-def transform(col: ColumnOrName, f: Callable[[Column], Column]) -> Column: ...
-@overload
-def transform(col: ColumnOrName, f: Callable[[Column, Column], Column]) ->
Column: ...
-def exists(col: ColumnOrName, f: Callable[[Column], Column]) -> Column: ...
-def forall(col: ColumnOrName, f: Callable[[Column], Column]) -> Column: ...
-@overload
-def filter(col: ColumnOrName, f: Callable[[Column], Column]) -> Column: ...
-@overload
-def filter(col: ColumnOrName, f: Callable[[Column, Column], Column]) ->
Column: ...
-def aggregate(
- col: ColumnOrName,
- initialValue: ColumnOrName,
- merge: Callable[[Column, Column], Column],
- finish: Optional[Callable[[Column], Column]] = ...,
-) -> Column: ...
-def zip_with(
- left: ColumnOrName,
- right: ColumnOrName,
- f: Callable[[Column, Column], Column],
-) -> Column: ...
-def transform_keys(
- col: ColumnOrName, f: Callable[[Column, Column], Column]
-) -> Column: ...
-def transform_values(
- col: ColumnOrName, f: Callable[[Column, Column], Column]
-) -> Column: ...
-def map_filter(col: ColumnOrName, f: Callable[[Column, Column], Column]) ->
Column: ...
-def map_zip_with(
- col1: ColumnOrName,
- col2: ColumnOrName,
- f: Callable[[Column, Column, Column], Column],
-) -> Column: ...
-def abs(col: ColumnOrName) -> Column: ...
-def acos(col: ColumnOrName) -> Column: ...
-def acosh(col: ColumnOrName) -> Column: ...
-def asc(col: ColumnOrName) -> Column: ...
-def asc_nulls_first(col: ColumnOrName) -> Column: ...
-def asc_nulls_last(col: ColumnOrName) -> Column: ...
-def ascii(col: ColumnOrName) -> Column: ...
-def asin(col: ColumnOrName) -> Column: ...
-def asinh(col: ColumnOrName) -> Column: ...
-def atan(col: ColumnOrName) -> Column: ...
-def atanh(col: ColumnOrName) -> Column: ...
-@overload
-def atan2(col1: ColumnOrName, col2: ColumnOrName) -> Column: ...
-@overload
-def atan2(col1: float, col2: ColumnOrName) -> Column: ...
-@overload
-def atan2(col1: ColumnOrName, col2: float) -> Column: ...
-def avg(col: ColumnOrName) -> Column: ...
-def base64(col: ColumnOrName) -> Column: ...
-def bitwiseNOT(col: ColumnOrName) -> Column: ...
-def bitwise_not(col: ColumnOrName) -> Column: ...
-def cbrt(col: ColumnOrName) -> Column: ...
-def ceil(col: ColumnOrName) -> Column: ...
-def col(col: str) -> Column: ...
-def collect_list(col: ColumnOrName) -> Column: ...
-def collect_set(col: ColumnOrName) -> Column: ...
-def column(col: str) -> Column: ...
-def cos(col: ColumnOrName) -> Column: ...
-def cosh(col: ColumnOrName) -> Column: ...
-def cot(col: ColumnOrName) -> Column: ...
-def count(col: ColumnOrName) -> Column: ...
-def csc(col: ColumnOrName) -> Column: ...
-def cume_dist() -> Column: ...
-def degrees(col: ColumnOrName) -> Column: ...
-def dense_rank() -> Column: ...
-def desc(col: ColumnOrName) -> Column: ...
-def desc_nulls_first(col: ColumnOrName) -> Column: ...
-def desc_nulls_last(col: ColumnOrName) -> Column: ...
-def exp(col: ColumnOrName) -> Column: ...
-def expm1(col: ColumnOrName) -> Column: ...
-def floor(col: ColumnOrName) -> Column: ...
-@overload
-def hypot(col1: ColumnOrName, col2: ColumnOrName) -> Column: ...
-@overload
-def hypot(col1: float, col2: ColumnOrName) -> Column: ...
-@overload
-def hypot(col1: ColumnOrName, col2: float) -> Column: ...
-def kurtosis(col: ColumnOrName) -> Column: ...
-def lit(col: Any) -> Column: ...
-def log10(col: ColumnOrName) -> Column: ...
-def log1p(col: ColumnOrName) -> Column: ...
-def lower(col: ColumnOrName) -> Column: ...
-def ltrim(col: ColumnOrName) -> Column: ...
-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: float, col2: ColumnOrName) -> Column: ...
-@overload
-def pow(col1: ColumnOrName, col2: float) -> Column: ...
-def product(col: ColumnOrName) -> Column: ...
-def radians(col: ColumnOrName) -> Column: ...
-def rank() -> Column: ...
-def rint(col: ColumnOrName) -> Column: ...
-def row_number() -> Column: ...
-def rtrim(col: ColumnOrName) -> Column: ...
-def sec(col: ColumnOrName) -> Column: ...
-def signum(col: ColumnOrName) -> Column: ...
-def sin(col: ColumnOrName) -> Column: ...
-def sinh(col: ColumnOrName) -> Column: ...
-def skewness(col: ColumnOrName) -> Column: ...
-def sqrt(col: ColumnOrName) -> Column: ...
-def stddev(col: ColumnOrName) -> Column: ...
-def stddev_pop(col: ColumnOrName) -> Column: ...
-def stddev_samp(col: ColumnOrName) -> Column: ...
-def sum(col: ColumnOrName) -> Column: ...
-def sumDistinct(col: ColumnOrName) -> Column: ...
-def sum_distinct(col: ColumnOrName) -> Column: ...
-def tan(col: ColumnOrName) -> Column: ...
-def tanh(col: ColumnOrName) -> Column: ...
-def toDegrees(col: ColumnOrName) -> Column: ...
-def toRadians(col: ColumnOrName) -> Column: ...
-def trim(col: ColumnOrName) -> Column: ...
-def unbase64(col: ColumnOrName) -> Column: ...
-def upper(col: ColumnOrName) -> Column: ...
-def var_pop(col: ColumnOrName) -> Column: ...
-def var_samp(col: ColumnOrName) -> Column: ...
-def variance(col: ColumnOrName) -> Column: ...
-@overload
-def udf(
- f: Callable[..., Any], returnType: DataTypeOrString = ...
-) -> UserDefinedFunctionLike: ...
-@overload
-def udf(
- f: DataTypeOrString = ...,
-) -> Callable[[Callable[..., Any]], UserDefinedFunctionLike]: ...
-@overload
-def udf(
- *,
- returnType: DataTypeOrString = ...,
-) -> Callable[[Callable[..., Any]], UserDefinedFunctionLike]: ...
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
