You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ru...@apache.org on 2022/12/05 06:01:52 UTC
[spark] branch master updated: [SPARK-41380][CONNECT][PYTHON] Implement aggregation functions
This is an automated email from the ASF dual-hosted git repository.
ruifengz 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 a383dde45d6 [SPARK-41380][CONNECT][PYTHON] Implement aggregation functions
a383dde45d6 is described below
commit a383dde45d6b9983575438da016d122ab2f9873f
Author: Ruifeng Zheng <ru...@apache.org>
AuthorDate: Mon Dec 5 14:01:30 2022 +0800
[SPARK-41380][CONNECT][PYTHON] Implement aggregation functions
### What changes were proposed in this pull request?
Implement [aggregation functions](https://github.com/apache/spark/blob/master/python/docs/source/reference/pyspark.sql/functions.rst#aggregate-functions), except:
1. `approxCountDistinct`, `countDistinct`, `sumDistinct ` - deprecated
2. `count_distinct`, `sum_distinct` - need to support `isDistinct` in proto message `UnresolvedFunction`, I will do it in a followup;
3. `product` - need to add it in `FunctionRegistry`, I will do it in a followup;
also note that:
1. `grouping` and `grouping` are added in this PR, but can not be used for now: `DataFrame.{cube, rollup}` are still missing, and `grouping() can only be used with GroupingSets/Cube/Rollup`;
2. `median` and `mode` are new APIs added in 3.4.0;
### Why are the changes needed?
For API coverage
### Does this PR introduce _any_ user-facing change?
new API
### How was this patch tested?
added UT
Closes #38905 from zhengruifeng/connect_function_agg.
Authored-by: Ruifeng Zheng <ru...@apache.org>
Signed-off-by: Ruifeng Zheng <ru...@apache.org>
---
python/pyspark/sql/connect/functions.py | 1099 +++++++++++++++++++-
.../sql/tests/connect/test_connect_function.py | 90 ++
2 files changed, 1188 insertions(+), 1 deletion(-)
diff --git a/python/pyspark/sql/connect/functions.py b/python/pyspark/sql/connect/functions.py
index 9b1163a1bdf..a52fc58fd0a 100644
--- a/python/pyspark/sql/connect/functions.py
+++ b/python/pyspark/sql/connect/functions.py
@@ -23,7 +23,7 @@ from pyspark.sql.connect.column import (
SQLExpression,
)
-from typing import Any, TYPE_CHECKING, Union, List, Optional
+from typing import Any, TYPE_CHECKING, Union, List, Optional, Tuple
if TYPE_CHECKING:
from pyspark.sql.connect._typing import ColumnOrName
@@ -2112,3 +2112,1100 @@ def unhex(col: "ColumnOrName") -> Column:
[Row(unhex(a)=bytearray(b'ABC'))]
"""
return _invoke_function_over_columns("unhex", col)
+
+
+# Aggregate Functions
+
+
+# def approxCountDistinct(col: "ColumnOrName", rsd: Optional[float] = None) -> Column:
+# """
+# .. versionadded:: 1.3.0
+#
+# .. deprecated:: 2.1.0
+# Use :func:`approx_count_distinct` instead.
+# """
+# warnings.warn("Deprecated in 2.1, use approx_count_distinct instead.", FutureWarning)
+# return approx_count_distinct(col, rsd)
+
+
+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`.
+
+ .. versionadded:: 3.4.0
+
+ Parameters
+ ----------
+ col : :class:`~pyspark.sql.Column` or str
+ rsd : float, optional
+ maximum relative standard deviation allowed (default = 0.05).
+ For rsd < 0.01, it is more efficient to use :func:`count_distinct`
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ the column of computed results.
+
+ Examples
+ --------
+ >>> df = spark.createDataFrame([1,2,2,3], "INT")
+ >>> df.agg(approx_count_distinct("value").alias('distinct_values')).show()
+ +---------------+
+ |distinct_values|
+ +---------------+
+ | 3|
+ +---------------+
+ """
+ if rsd is None:
+ return _invoke_function("approx_count_distinct", _to_col(col))
+ else:
+ return _invoke_function("approx_count_distinct", _to_col(col), lit(rsd))
+
+
+def avg(col: "ColumnOrName") -> Column:
+ """
+ Aggregate function: returns the average of the values in a group.
+
+ .. versionadded:: 3.4.0
+
+ Parameters
+ ----------
+ col : :class:`~pyspark.sql.Column` or str
+ target column to compute on.
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ the column for computed results.
+
+ Examples
+ --------
+ >>> df = spark.range(10)
+ >>> df.select(avg(col("id"))).show()
+ +-------+
+ |avg(id)|
+ +-------+
+ | 4.5|
+ +-------+
+ """
+ return _invoke_function_over_columns("avg", col)
+
+
+def collect_list(col: "ColumnOrName") -> Column:
+ """
+ Aggregate function: returns a list of objects with duplicates.
+
+ .. versionadded:: 3.4.0
+
+ Notes
+ -----
+ The function is non-deterministic because the order of collected results depends
+ on the order of the rows which may be non-deterministic after a shuffle.
+
+ Parameters
+ ----------
+ col : :class:`~pyspark.sql.Column` or str
+ target column to compute on.
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ list of objects with duplicates.
+
+ Examples
+ --------
+ >>> df2 = spark.createDataFrame([(2,), (5,), (5,)], ('age',))
+ >>> df2.agg(collect_list('age')).collect()
+ [Row(collect_list(age)=[2, 5, 5])]
+ """
+ return _invoke_function_over_columns("collect_list", col)
+
+
+def collect_set(col: "ColumnOrName") -> Column:
+ """
+ Aggregate function: returns a set of objects with duplicate elements eliminated.
+
+ .. versionadded:: 3.4.0
+
+ Notes
+ -----
+ The function is non-deterministic because the order of collected results depends
+ on the order of the rows which may be non-deterministic after a shuffle.
+
+ Parameters
+ ----------
+ col : :class:`~pyspark.sql.Column` or str
+ target column to compute on.
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ list of objects with no duplicates.
+
+ Examples
+ --------
+ >>> df2 = spark.createDataFrame([(2,), (5,), (5,)], ('age',))
+ >>> df2.agg(array_sort(collect_set('age')).alias('c')).collect()
+ [Row(c=[2, 5])]
+ """
+ return _invoke_function_over_columns("collect_set", col)
+
+
+def corr(col1: "ColumnOrName", col2: "ColumnOrName") -> Column:
+ """Returns a new :class:`~pyspark.sql.Column` for the Pearson Correlation Coefficient for
+ ``col1`` and ``col2``.
+
+ .. versionadded:: 3.4.0
+
+ Parameters
+ ----------
+ col1 : :class:`~pyspark.sql.Column` or str
+ first column to calculate correlation.
+ col1 : :class:`~pyspark.sql.Column` or str
+ second column to calculate correlation.
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ Pearson Correlation Coefficient of these two column values.
+
+ Examples
+ --------
+ >>> a = range(20)
+ >>> b = [2 * x for x in range(20)]
+ >>> df = spark.createDataFrame(zip(a, b), ["a", "b"])
+ >>> df.agg(corr("a", "b").alias('c')).collect()
+ [Row(c=1.0)]
+ """
+ return _invoke_function_over_columns("corr", col1, col2)
+
+
+def count(col: "ColumnOrName") -> Column:
+ """
+ Aggregate function: returns the number of items in a group.
+
+ .. versionadded:: 3.4.0
+
+ Parameters
+ ----------
+ col : :class:`~pyspark.sql.Column` or str
+ target column to compute on.
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ column for computed results.
+
+ Examples
+ --------
+ Count by all columns (start), and by a column that does not count ``None``.
+
+ >>> df = spark.createDataFrame([(None,), ("a",), ("b",), ("c",)], schema=["alphabets"])
+ >>> df.select(count(expr("*")), count(df.alphabets)).show()
+ +--------+----------------+
+ |count(1)|count(alphabets)|
+ +--------+----------------+
+ | 4| 3|
+ +--------+----------------+
+ """
+ return _invoke_function_over_columns("count", col)
+
+
+# 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`
+# directly.
+#
+# .. versionadded:: 1.3.0
+# """
+# return count_distinct(col, *cols)
+
+
+# TODO(SPARK-41381): add isDistinct in UnresolvedFunction
+# def count_distinct(col: "ColumnOrName", *cols: "ColumnOrName") -> Column:
+# """Returns a new :class:`Column` for distinct count of ``col`` or ``cols``.
+#
+# .. versionadded:: 3.4.0
+#
+# Parameters
+# ----------
+# col : :class:`~pyspark.sql.Column` or str
+# first column to compute on.
+# cols : :class:`~pyspark.sql.Column` or str
+# other columns to compute on.
+#
+# Returns
+# -------
+# :class:`~pyspark.sql.Column`
+# distinct values of these two column values.
+#
+# Examples
+# --------
+# >>> from pyspark.sql import types
+# >>> df1 = spark.createDataFrame([1, 1, 3], types.IntegerType())
+# >>> df2 = spark.createDataFrame([1, 2], types.IntegerType())
+# >>> df1.join(df2).show()
+# +-----+-----+
+# |value|value|
+# +-----+-----+
+# | 1| 1|
+# | 1| 2|
+# | 1| 1|
+# | 1| 2|
+# | 3| 1|
+# | 3| 2|
+# +-----+-----+
+# >>> df1.join(df2).select(count_distinct(df1.value, df2.value)).show()
+# +----------------------------+
+# |count(DISTINCT value, value)|
+# +----------------------------+
+# | 4|
+# +----------------------------+
+# """
+# return _invoke_function(
+# "count_distinct", _to_java_column(col), _to_seq(sc, cols, _to_java_column)
+# )
+
+
+def covar_pop(col1: "ColumnOrName", col2: "ColumnOrName") -> Column:
+ """Returns a new :class:`~pyspark.sql.Column` for the population covariance of ``col1`` and
+ ``col2``.
+
+ .. versionadded:: 3.4.0
+
+ Parameters
+ ----------
+ col1 : :class:`~pyspark.sql.Column` or str
+ first column to calculate covariance.
+ col1 : :class:`~pyspark.sql.Column` or str
+ second column to calculate covariance.
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ covariance of these two column values.
+
+ Examples
+ --------
+ >>> a = [1] * 10
+ >>> b = [1] * 10
+ >>> df = spark.createDataFrame(zip(a, b), ["a", "b"])
+ >>> df.agg(covar_pop("a", "b").alias('c')).collect()
+ [Row(c=0.0)]
+ """
+ return _invoke_function_over_columns("covar_pop", 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``.
+
+ .. versionadded:: 3.4.0
+
+ Parameters
+ ----------
+ col1 : :class:`~pyspark.sql.Column` or str
+ first column to calculate covariance.
+ col1 : :class:`~pyspark.sql.Column` or str
+ second column to calculate covariance.
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ sample covariance of these two column values.
+
+ Examples
+ --------
+ >>> a = [1] * 10
+ >>> b = [1] * 10
+ >>> df = spark.createDataFrame(zip(a, b), ["a", "b"])
+ >>> df.agg(covar_samp("a", "b").alias('c')).collect()
+ [Row(c=0.0)]
+ """
+ return _invoke_function_over_columns("covar_samp", col1, col2)
+
+
+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
+ value it sees when ignoreNulls is set to true. If all values are null, then null is returned.
+
+ .. versionadded:: 3.4.0
+
+ Notes
+ -----
+ The function is non-deterministic because its results depends on the order of the
+ rows which may be non-deterministic after a shuffle.
+
+ Parameters
+ ----------
+ col : :class:`~pyspark.sql.Column` or str
+ column to fetch first value for.
+ ignorenulls : :class:`~pyspark.sql.Column` or str
+ if first value is null then look for first non-null value.
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ first value of the group.
+
+ Examples
+ --------
+ >>> df = spark.createDataFrame([("Alice", 2), ("Bob", 5), ("Alice", None)], ("name", "age"))
+ >>> df = df.orderBy(df.age)
+ >>> df.groupby("name").agg(first("age")).orderBy("name").show()
+ +-----+----------+
+ | name|first(age)|
+ +-----+----------+
+ |Alice| null|
+ | Bob| 5|
+ +-----+----------+
+
+ Now, to ignore any nulls we needs to set ``ignorenulls`` to `True`
+
+ >>> df.groupby("name").agg(first("age", ignorenulls=True)).orderBy("name").show()
+ +-----+----------+
+ | name|first(age)|
+ +-----+----------+
+ |Alice| 2|
+ | Bob| 5|
+ +-----+----------+
+ """
+ return _invoke_function("first", _to_col(col), lit(ignorenulls))
+
+
+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.
+
+ .. versionadded:: 3.4.0
+
+ Parameters
+ ----------
+ col : :class:`~pyspark.sql.Column` or str
+ column to check if it's aggregated.
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ returns 1 for aggregated or 0 for not aggregated in the result set.
+
+ Examples
+ --------
+ >>> df = spark.createDataFrame([("Alice", 2), ("Bob", 5)], ("name", "age"))
+ >>> df.cube("name").agg(grouping("name"), sum("age")).orderBy("name").show()
+ +-----+--------------+--------+
+ | name|grouping(name)|sum(age)|
+ +-----+--------------+--------+
+ | null| 1| 7|
+ |Alice| 0| 2|
+ | Bob| 0| 5|
+ +-----+--------------+--------+
+ """
+ return _invoke_function_over_columns("grouping", col)
+
+
+def grouping_id(*cols: "ColumnOrName") -> Column:
+ """
+ Aggregate function: returns the level of grouping, equals to
+
+ (grouping(c1) << (n-1)) + (grouping(c2) << (n-2)) + ... + grouping(cn)
+
+ .. versionadded:: 3.4.0
+
+ Notes
+ -----
+ The list of columns should match with grouping columns exactly, or empty (means all
+ the grouping columns).
+
+ Parameters
+ ----------
+ cols : :class:`~pyspark.sql.Column` or str
+ columns to check for.
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ returns level of the grouping it relates to.
+
+ Examples
+ --------
+ >>> df = spark.createDataFrame([(1, "a", "a"),
+ ... (3, "a", "a"),
+ ... (4, "b", "c")], ["c1", "c2", "c3"])
+ >>> df.cube("c2", "c3").agg(grouping_id(), sum("c1")).orderBy("c2", "c3").show()
+ +----+----+-------------+-------+
+ | c2| c3|grouping_id()|sum(c1)|
+ +----+----+-------------+-------+
+ |null|null| 3| 8|
+ |null| a| 2| 4|
+ |null| c| 2| 4|
+ | a|null| 1| 4|
+ | a| a| 0| 4|
+ | b|null| 1| 4|
+ | b| c| 0| 4|
+ +----+----+-------------+-------+
+ """
+ return _invoke_function_over_columns("grouping_id", *cols)
+
+
+def kurtosis(col: "ColumnOrName") -> Column:
+ """
+ Aggregate function: returns the kurtosis of the values in a group.
+
+ .. versionadded:: 3.4.0
+
+ Parameters
+ ----------
+ col : :class:`~pyspark.sql.Column` or str
+ target column to compute on.
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ kurtosis of given column.
+
+ Examples
+ --------
+ >>> df = spark.createDataFrame([[1],[1],[2]], ["c"])
+ >>> df.select(kurtosis(df.c)).show()
+ +-----------+
+ |kurtosis(c)|
+ +-----------+
+ | -1.5|
+ +-----------+
+ """
+ return _invoke_function_over_columns("kurtosis", col)
+
+
+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
+ value it sees when ignoreNulls is set to true. If all values are null, then null is returned.
+
+ .. versionadded:: 3.4.0
+
+ Notes
+ -----
+ The function is non-deterministic because its results depends on the order of the
+ rows which may be non-deterministic after a shuffle.
+
+ Parameters
+ ----------
+ col : :class:`~pyspark.sql.Column` or str
+ column to fetch last value for.
+ ignorenulls : :class:`~pyspark.sql.Column` or str
+ if last value is null then look for non-null value.
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ last value of the group.
+
+ Examples
+ --------
+ >>> df = spark.createDataFrame([("Alice", 2), ("Bob", 5), ("Alice", None)], ("name", "age"))
+ >>> df = df.orderBy(df.age.desc())
+ >>> df.groupby("name").agg(last("age")).orderBy("name").show()
+ +-----+---------+
+ | name|last(age)|
+ +-----+---------+
+ |Alice| null|
+ | Bob| 5|
+ +-----+---------+
+
+ Now, to ignore any nulls we needs to set ``ignorenulls`` to `True`
+
+ >>> df.groupby("name").agg(last("age", ignorenulls=True)).orderBy("name").show()
+ +-----+---------+
+ | name|last(age)|
+ +-----+---------+
+ |Alice| 2|
+ | Bob| 5|
+ +-----+---------+
+ """
+ return _invoke_function("last", _to_col(col), lit(ignorenulls))
+
+
+def max(col: "ColumnOrName") -> Column:
+ """
+ Aggregate function: returns the maximum value of the expression in a group.
+
+ .. versionadded:: 3.4.0
+
+ Parameters
+ ----------
+ col : :class:`~pyspark.sql.Column` or str
+ target column to compute on.
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ column for computed results.
+
+ Examples
+ --------
+ >>> df = spark.range(10)
+ >>> df.select(max(col("id"))).show()
+ +-------+
+ |max(id)|
+ +-------+
+ | 9|
+ +-------+
+ """
+ return _invoke_function_over_columns("max", col)
+
+
+def max_by(col: "ColumnOrName", ord: "ColumnOrName") -> Column:
+ """
+ Returns the value associated with the maximum value of ord.
+
+ .. versionadded:: 3.4.0
+
+ Parameters
+ ----------
+ col : :class:`~pyspark.sql.Column` or str
+ target column to compute on.
+ ord : :class:`~pyspark.sql.Column` or str
+ column to be maximized
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ value associated with the maximum value of ord.
+
+ Examples
+ --------
+ >>> df = spark.createDataFrame([
+ ... ("Java", 2012, 20000), ("dotNET", 2012, 5000),
+ ... ("dotNET", 2013, 48000), ("Java", 2013, 30000)],
+ ... schema=("course", "year", "earnings"))
+ >>> df.groupby("course").agg(max_by("year", "earnings")).show()
+ +------+----------------------+
+ |course|max_by(year, earnings)|
+ +------+----------------------+
+ | Java| 2013|
+ |dotNET| 2013|
+ +------+----------------------+
+ """
+ return _invoke_function_over_columns("max_by", col, ord)
+
+
+def mean(col: "ColumnOrName") -> Column:
+ """
+ Aggregate function: returns the average of the values in a group.
+ An alias of :func:`avg`.
+
+ .. versionadded:: 1.4.0
+
+ Parameters
+ ----------
+ col : :class:`~pyspark.sql.Column` or str
+ target column to compute on.
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ the column for computed results.
+
+ Examples
+ --------
+ >>> df = spark.range(10)
+ >>> df.select(mean(df.id)).show()
+ +-------+
+ |avg(id)|
+ +-------+
+ | 4.5|
+ +-------+
+ """
+ return avg(col)
+
+
+def median(col: "ColumnOrName") -> Column:
+ """
+ Returns the median of the values in a group.
+
+ .. versionadded:: 3.4.0
+
+ Parameters
+ ----------
+ col : :class:`~pyspark.sql.Column` or str
+ target column to compute on.
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ the median of the values in a group.
+
+ Examples
+ --------
+ >>> df = spark.createDataFrame([
+ ... ("Java", 2012, 20000), ("dotNET", 2012, 5000),
+ ... ("Java", 2012, 22000), ("dotNET", 2012, 10000),
+ ... ("dotNET", 2013, 48000), ("Java", 2013, 30000)],
+ ... schema=("course", "year", "earnings"))
+ >>> df.groupby("course").agg(median("earnings")).show()
+ +------+----------------+
+ |course|median(earnings)|
+ +------+----------------+
+ | Java| 22000.0|
+ |dotNET| 10000.0|
+ +------+----------------+
+ """
+ return _invoke_function_over_columns("median", col)
+
+
+def min(col: "ColumnOrName") -> Column:
+ """
+ Aggregate function: returns the minimum value of the expression in a group.
+
+ .. versionadded:: 3.4.0
+
+ Parameters
+ ----------
+ col : :class:`~pyspark.sql.Column` or str
+ target column to compute on.
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ column for computed results.
+
+ Examples
+ --------
+ >>> df = spark.range(10)
+ >>> df.select(min(df.id)).show()
+ +-------+
+ |min(id)|
+ +-------+
+ | 0|
+ +-------+
+ """
+ return _invoke_function_over_columns("min", col)
+
+
+def min_by(col: "ColumnOrName", ord: "ColumnOrName") -> Column:
+ """
+ Returns the value associated with the minimum value of ord.
+
+ .. versionadded:: 3.4.0
+
+ Parameters
+ ----------
+ col : :class:`~pyspark.sql.Column` or str
+ target column to compute on.
+ ord : :class:`~pyspark.sql.Column` or str
+ column to be minimized
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ value associated with the minimum value of ord.
+
+ Examples
+ --------
+ >>> df = spark.createDataFrame([
+ ... ("Java", 2012, 20000), ("dotNET", 2012, 5000),
+ ... ("dotNET", 2013, 48000), ("Java", 2013, 30000)],
+ ... schema=("course", "year", "earnings"))
+ >>> df.groupby("course").agg(min_by("year", "earnings")).show()
+ +------+----------------------+
+ |course|min_by(year, earnings)|
+ +------+----------------------+
+ | Java| 2012|
+ |dotNET| 2012|
+ +------+----------------------+
+ """
+ return _invoke_function_over_columns("min_by", col, ord)
+
+
+def mode(col: "ColumnOrName") -> Column:
+ """
+ Returns the most frequent value in a group.
+
+ .. versionadded:: 3.4.0
+
+ Parameters
+ ----------
+ col : :class:`~pyspark.sql.Column` or str
+ target column to compute on.
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ the most frequent value in a group.
+
+ Examples
+ --------
+ >>> df = spark.createDataFrame([
+ ... ("Java", 2012, 20000), ("dotNET", 2012, 5000),
+ ... ("Java", 2012, 20000), ("dotNET", 2012, 5000),
+ ... ("dotNET", 2013, 48000), ("Java", 2013, 30000)],
+ ... schema=("course", "year", "earnings"))
+ >>> df.groupby("course").agg(mode("year")).show()
+ +------+----------+
+ |course|mode(year)|
+ +------+----------+
+ | Java| 2012|
+ |dotNET| 2012|
+ +------+----------+
+ """
+ return _invoke_function_over_columns("mode", col)
+
+
+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.
+
+
+ .. versionadded:: 3.4.0
+
+ Parameters
+ ----------
+ col : :class:`~pyspark.sql.Column` or str
+ input column.
+ percentage : :class:`~pyspark.sql.Column`, float, list of floats or tuple of floats
+ percentage in decimal (must be between 0.0 and 1.0).
+ When percentage is an array, each value of the percentage array must be between 0.0 and 1.0.
+ In this case, returns the approximate percentile array of column col
+ at the given percentage array.
+ accuracy : :class:`~pyspark.sql.Column` or float
+ is a positive numeric literal which controls approximation accuracy
+ at the cost of memory. Higher value of accuracy yields better accuracy,
+ 1.0/accuracy is the relative error of the approximation. (default: 10000).
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ approximate `percentile` of the numeric column.
+
+ Examples
+ --------
+ >>> key = (col("id") % 3).alias("key")
+ >>> value = (randn(42) + key * 10).alias("value")
+ >>> df = spark.range(0, 1000, 1, 1).select(key, value)
+ >>> df.select(
+ ... percentile_approx("value", [0.25, 0.5, 0.75], 1000000).alias("quantiles")
+ ... ).printSchema()
+ root
+ |-- quantiles: array (nullable = true)
+ | |-- element: double (containsNull = false)
+
+ >>> df.groupBy("key").agg(
+ ... percentile_approx("value", 0.5, lit(1000000)).alias("median")
+ ... ).printSchema()
+ root
+ |-- key: long (nullable = true)
+ |-- median: double (nullable = true)
+ """
+
+ if isinstance(percentage, Column):
+ percentage_col = percentage
+ elif isinstance(percentage, (list, tuple)):
+ # Convert tuple to list
+ percentage_col = lit(list(percentage))
+ else:
+ # Probably scalar
+ percentage_col = lit(percentage)
+
+ return _invoke_function("percentile_approx", _to_col(col), percentage_col, lit(accuracy))
+
+
+# TODO(SPARK-41382): add product in FunctionRegistry?
+# def product(col: "ColumnOrName") -> Column:
+# """
+# Aggregate function: returns the product of the values in a group.
+#
+# .. versionadded:: 3.4.0
+#
+# Parameters
+# ----------
+# col : str, :class:`Column`
+# column containing values to be multiplied together
+#
+# Returns
+# -------
+# :class:`~pyspark.sql.Column`
+# the column for computed results.
+#
+# Examples
+# --------
+# >>> df = spark.range(1, 10).toDF('x').withColumn('mod3', col('x') % 3)
+# >>> prods = df.groupBy('mod3').agg(product('x').alias('product'))
+# >>> prods.orderBy('mod3').show()
+# +----+-------+
+# |mod3|product|
+# +----+-------+
+# | 0| 162.0|
+# | 1| 28.0|
+# | 2| 80.0|
+# +----+-------+
+# """
+# return _invoke_function_over_columns("product", col)
+
+
+def skewness(col: "ColumnOrName") -> Column:
+ """
+ Aggregate function: returns the skewness of the values in a group.
+
+ .. versionadded:: 3.4.0
+
+ Parameters
+ ----------
+ col : :class:`~pyspark.sql.Column` or str
+ target column to compute on.
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ skewness of given column.
+
+ Examples
+ --------
+ >>> df = spark.createDataFrame([[1],[1],[2]], ["c"])
+ >>> df.select(skewness(df.c)).first()
+ Row(skewness(c)=0.70710...)
+ """
+ return _invoke_function_over_columns("skewness", col)
+
+
+def stddev(col: "ColumnOrName") -> Column:
+ """
+ Aggregate function: alias for stddev_samp.
+
+ .. versionadded:: 3.4.0
+
+ Parameters
+ ----------
+ col : :class:`~pyspark.sql.Column` or str
+ target column to compute on.
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ standard deviation of given column.
+
+ Examples
+ --------
+ >>> df = spark.range(6)
+ >>> df.select(stddev(df.id)).first()
+ Row(stddev_samp(id)=1.87082...)
+ """
+ return stddev_samp(col)
+
+
+def stddev_samp(col: "ColumnOrName") -> Column:
+ """
+ Aggregate function: returns the unbiased sample standard deviation of
+ the expression in a group.
+
+ .. versionadded:: 3.4.0
+
+ Parameters
+ ----------
+ col : :class:`~pyspark.sql.Column` or str
+ target column to compute on.
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ standard deviation of given column.
+
+ Examples
+ --------
+ >>> df = spark.range(6)
+ >>> df.select(stddev_samp(df.id)).first()
+ Row(stddev_samp(id)=1.87082...)
+ """
+ return _invoke_function_over_columns("stddev_samp", col)
+
+
+def stddev_pop(col: "ColumnOrName") -> Column:
+ """
+ Aggregate function: returns population standard deviation of
+ the expression in a group.
+
+ .. versionadded:: 3.4.0
+
+ Parameters
+ ----------
+ col : :class:`~pyspark.sql.Column` or str
+ target column to compute on.
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ standard deviation of given column.
+
+ Examples
+ --------
+ >>> df = spark.range(6)
+ >>> df.select(stddev_pop(df.id)).first()
+ Row(stddev_pop(id)=1.70782...)
+ """
+ return _invoke_function_over_columns("stddev_pop", col)
+
+
+def sum(col: "ColumnOrName") -> Column:
+ """
+ Aggregate function: returns the sum of all values in the expression.
+
+ .. versionadded:: 3.4.0
+
+ Parameters
+ ----------
+ col : :class:`~pyspark.sql.Column` or str
+ target column to compute on.
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ the column for computed results.
+
+ Examples
+ --------
+ >>> df = spark.range(10)
+ >>> df.select(sum(df["id"])).show()
+ +-------+
+ |sum(id)|
+ +-------+
+ | 45|
+ +-------+
+ """
+ return _invoke_function_over_columns("sum", col)
+
+
+# def sumDistinct(col: "ColumnOrName") -> Column:
+# """
+# Aggregate function: returns the sum of distinct values in the expression.
+#
+# .. versionadded:: 1.3.0
+#
+# .. deprecated:: 3.2.0
+# Use :func:`sum_distinct` instead.
+# """
+# warnings.warn("Deprecated in 3.2, use sum_distinct instead.", FutureWarning)
+# return sum_distinct(col)
+
+
+# TODO(SPARK-41381): add isDistinct in UnresolvedFunction
+# def sum_distinct(col: "ColumnOrName") -> Column:
+# """
+# Aggregate function: returns the sum of distinct values in the expression.
+#
+# .. versionadded:: 3.4.0
+#
+# Parameters
+# ----------
+# col : :class:`~pyspark.sql.Column` or str
+# target column to compute on.
+#
+# Returns
+# -------
+# :class:`~pyspark.sql.Column`
+# the column for computed results.
+#
+# Examples
+# --------
+# >>> df = spark.createDataFrame([(None,), (1,), (1,), (2,)], schema=["numbers"])
+# >>> df.select(sum_distinct(col("numbers"))).show()
+# +---------------------+
+# |sum(DISTINCT numbers)|
+# +---------------------+
+# | 3|
+# +---------------------+
+# """
+# return _invoke_function_over_columns("sum_distinct", col)
+
+
+def var_pop(col: "ColumnOrName") -> Column:
+ """
+ Aggregate function: returns the population variance of the values in a group.
+
+ .. versionadded:: 3.4.0
+
+ Parameters
+ ----------
+ col : :class:`~pyspark.sql.Column` or str
+ target column to compute on.
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ variance of given column.
+
+ Examples
+ --------
+ >>> df = spark.range(6)
+ >>> df.select(var_pop(df.id)).first()
+ Row(var_pop(id)=2.91666...)
+ """
+ return _invoke_function_over_columns("var_pop", col)
+
+
+def var_samp(col: "ColumnOrName") -> Column:
+ """
+ Aggregate function: returns the unbiased sample variance of
+ the values in a group.
+
+ .. versionadded:: 3.4.0
+
+ Parameters
+ ----------
+ col : :class:`~pyspark.sql.Column` or str
+ target column to compute on.
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ variance of given column.
+
+ Examples
+ --------
+ >>> df = spark.range(6)
+ >>> df.select(var_samp(df.id)).show()
+ +------------+
+ |var_samp(id)|
+ +------------+
+ | 3.5|
+ +------------+
+ """
+ return _invoke_function_over_columns("var_samp", col)
+
+
+def variance(col: "ColumnOrName") -> Column:
+ """
+ Aggregate function: alias for var_samp
+
+ .. versionadded:: 3.4.0
+
+ Parameters
+ ----------
+ col : :class:`~pyspark.sql.Column` or str
+ target column to compute on.
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ variance of given column.
+
+ Examples
+ --------
+ >>> df = spark.range(6)
+ >>> df.select(variance(df.id)).show()
+ +------------+
+ |var_samp(id)|
+ +------------+
+ | 3.5|
+ +------------+
+ """
+ return var_samp(col)
diff --git a/python/pyspark/sql/tests/connect/test_connect_function.py b/python/pyspark/sql/tests/connect/test_connect_function.py
index 41b404edcab..83d27235bdb 100644
--- a/python/pyspark/sql/tests/connect/test_connect_function.py
+++ b/python/pyspark/sql/tests/connect/test_connect_function.py
@@ -320,6 +320,96 @@ class SparkConnectFunctionTests(SparkConnectFuncTestCase):
sdf.select(SF.shiftrightunsigned("b", 1)).toPandas(),
)
+ def test_aggregation_functions(self):
+ from pyspark.sql import functions as SF
+ from pyspark.sql.connect import functions as CF
+
+ query = """
+ SELECT * FROM VALUES
+ (0, float("NAN"), NULL), (1, NULL, 2.0), (1, 2.1, 3.5), (0, 0.5, 1.0)
+ AS tab(a, b, c)
+ """
+ # +---+----+----+
+ # | a| b| c|
+ # +---+----+----+
+ # | 0| NaN|null|
+ # | 1|null| 2.0|
+ # | 1| 2.1| 3.5|
+ # | 0| 0.5| 1.0|
+ # +---+----+----+
+
+ cdf = self.connect.sql(query)
+ sdf = self.spark.sql(query)
+
+ # TODO(SPARK-41383): add tests for grouping, grouping_id after DataFrame.cube is supported.
+ for cfunc, sfunc in [
+ (CF.approx_count_distinct, SF.approx_count_distinct),
+ (CF.avg, SF.avg),
+ (CF.collect_list, SF.collect_list),
+ (CF.collect_set, SF.collect_set),
+ (CF.count, SF.count),
+ # (CF.count_distinct, SF.count_distinct),
+ (CF.first, SF.first),
+ (CF.kurtosis, SF.kurtosis),
+ (CF.last, SF.last),
+ (CF.max, SF.max),
+ (CF.mean, SF.mean),
+ (CF.median, SF.median),
+ (CF.min, SF.min),
+ (CF.mode, SF.mode),
+ (CF.skewness, SF.skewness),
+ (CF.stddev, SF.stddev),
+ (CF.stddev_pop, SF.stddev_pop),
+ (CF.stddev_samp, SF.stddev_samp),
+ (CF.sum, SF.sum),
+ # (CF.sum_distinct, SF.sum_distinct),
+ (CF.var_pop, SF.var_pop),
+ (CF.var_samp, SF.var_samp),
+ (CF.variance, SF.variance),
+ ]:
+ self.assert_eq(
+ cdf.select(cfunc("b"), cfunc(cdf.c)).toPandas(),
+ sdf.select(sfunc("b"), sfunc(sdf.c)).toPandas(),
+ )
+ self.assert_eq(
+ cdf.groupBy("a").agg([cfunc("b"), cfunc(cdf.c)]).toPandas(),
+ sdf.groupBy("a").agg(sfunc("b"), sfunc(sdf.c)).toPandas(),
+ )
+
+ for cfunc, sfunc in [
+ (CF.corr, SF.corr),
+ (CF.covar_pop, SF.covar_pop),
+ (CF.covar_samp, SF.covar_samp),
+ (CF.max_by, SF.max_by),
+ (CF.min_by, SF.min_by),
+ ]:
+ self.assert_eq(
+ cdf.select(cfunc(cdf.b, "c")).toPandas(),
+ sdf.select(sfunc(sdf.b, "c")).toPandas(),
+ )
+ self.assert_eq(
+ cdf.groupBy("a").agg([cfunc(cdf.b, "c")]).toPandas(),
+ sdf.groupBy("a").agg(sfunc(sdf.b, "c")).toPandas(),
+ )
+
+ # test percentile_approx
+ self.assert_eq(
+ cdf.select(CF.percentile_approx(cdf.b, 0.5, 1000)).toPandas(),
+ sdf.select(SF.percentile_approx(sdf.b, 0.5, 1000)).toPandas(),
+ )
+ self.assert_eq(
+ cdf.select(CF.percentile_approx(cdf.b, [0.1, 0.9])).toPandas(),
+ sdf.select(SF.percentile_approx(sdf.b, [0.1, 0.9])).toPandas(),
+ )
+ self.assert_eq(
+ cdf.groupBy("a").agg([CF.percentile_approx("b", 0.5)]).toPandas(),
+ sdf.groupBy("a").agg(SF.percentile_approx("b", 0.5)).toPandas(),
+ )
+ self.assert_eq(
+ cdf.groupBy("a").agg([CF.percentile_approx(cdf.b, [0.1, 0.9])]).toPandas(),
+ sdf.groupBy("a").agg(SF.percentile_approx(sdf.b, [0.1, 0.9])).toPandas(),
+ )
+
if __name__ == "__main__":
from pyspark.sql.tests.connect.test_connect_function import * # noqa: F401
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org