| # |
| # 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 pyspark.sql.connect.utils import check_dependencies |
| |
| check_dependencies(__name__) |
| |
| import warnings |
| from typing import ( |
| Dict, |
| List, |
| Sequence, |
| Union, |
| TYPE_CHECKING, |
| Optional, |
| overload, |
| cast, |
| ) |
| |
| from pyspark.util import PythonEvalType |
| from pyspark.sql.group import GroupedData as PySparkGroupedData |
| from pyspark.sql.pandas.group_ops import PandasCogroupedOps as PySparkPandasCogroupedOps |
| from pyspark.sql.pandas.functions import _validate_vectorized_udf # type: ignore[attr-defined] |
| from pyspark.sql.pandas.typehints import infer_group_arrow_eval_type_from_func |
| from pyspark.sql.types import NumericType, StructType |
| |
| import pyspark.sql.connect.plan as plan |
| from pyspark.sql.column import Column |
| from pyspark.sql.connect.functions import builtin as F |
| from pyspark.errors import PySparkNotImplementedError, PySparkTypeError |
| from pyspark.sql.streaming.stateful_processor import StatefulProcessor |
| |
| if TYPE_CHECKING: |
| from pyspark.sql.connect._typing import ( |
| LiteralType, |
| PandasGroupedMapFunction, |
| GroupedMapPandasUserDefinedFunction, |
| PandasCogroupedMapFunction, |
| ArrowCogroupedMapFunction, |
| ArrowGroupedMapFunction, |
| PandasGroupedMapFunctionWithState, |
| ) |
| from pyspark.sql.connect.dataframe import DataFrame |
| from pyspark.sql.types import StructType |
| |
| |
| class GroupedData: |
| def __init__( |
| self, |
| df: "DataFrame", |
| group_type: str, |
| grouping_cols: Sequence[Column], |
| pivot_col: Optional[Column] = None, |
| pivot_values: Optional[Sequence["LiteralType"]] = None, |
| grouping_sets: Optional[Sequence[Sequence[Column]]] = None, |
| ) -> None: |
| from pyspark.sql.connect.dataframe import DataFrame |
| |
| assert isinstance(df, DataFrame) |
| self._df = df |
| |
| assert isinstance(group_type, str) and group_type in [ |
| "groupby", |
| "rollup", |
| "cube", |
| "pivot", |
| "grouping_sets", |
| ] |
| self._group_type = group_type |
| |
| assert isinstance(grouping_cols, list) and all(isinstance(g, Column) for g in grouping_cols) |
| self._grouping_cols: List[Column] = grouping_cols |
| |
| self._pivot_col: Optional["Column"] = None |
| self._pivot_values: Optional[List["Column"]] = None |
| if group_type == "pivot": |
| assert pivot_col is not None and isinstance(pivot_col, Column) |
| self._pivot_col = pivot_col |
| |
| if pivot_values is not None: |
| assert isinstance(pivot_values, list) |
| self._pivot_values = [F.lit(v) for v in pivot_values] |
| |
| self._grouping_sets: Optional[Sequence[Sequence["Column"]]] = None |
| if group_type == "grouping_sets": |
| assert grouping_sets is None or isinstance(grouping_sets, list) |
| self._grouping_sets = grouping_sets |
| |
| def __repr__(self) -> str: |
| # the expressions are not resolved here, |
| # so the string representation can be different from classic PySpark. |
| grouping_str = ", ".join(str(e._expr) for e in self._grouping_cols) |
| grouping_str = f"grouping expressions: [{grouping_str}]" |
| |
| value_str = ", ".join("%s: %s" % c for c in self._df.dtypes) |
| |
| if self._group_type == "groupby": |
| type_str = "GroupBy" |
| elif self._group_type == "rollup": |
| type_str = "RollUp" |
| elif self._group_type == "cube": |
| type_str = "Cube" |
| elif self._group_type == "grouping_sets": |
| type_str = "GroupingSets" |
| else: |
| type_str = "Pivot" |
| |
| return f"GroupedData[{grouping_str}, value: [{value_str}], type: {type_str}]" |
| |
| @overload |
| def agg(self, *exprs: Column) -> "DataFrame": |
| ... |
| |
| @overload |
| def agg(self, __exprs: Dict[str, str]) -> "DataFrame": |
| ... |
| |
| def agg(self, *exprs: Union[Column, Dict[str, str]]) -> "DataFrame": |
| from pyspark.sql.connect.dataframe import DataFrame |
| |
| assert exprs, "exprs should not be empty" |
| if len(exprs) == 1 and isinstance(exprs[0], dict): |
| # Convert the dict into key value pairs |
| aggregate_cols = [F._invoke_function(exprs[0][k], F.col(k)) for k in exprs[0]] |
| else: |
| # Columns |
| assert all(isinstance(c, Column) for c in exprs), "all exprs should be Column" |
| aggregate_cols = cast(List[Column], list(exprs)) |
| |
| return DataFrame( |
| plan.Aggregate( |
| child=self._df._plan, |
| group_type=self._group_type, |
| grouping_cols=self._grouping_cols, |
| aggregate_cols=aggregate_cols, |
| pivot_col=self._pivot_col, |
| pivot_values=self._pivot_values, |
| grouping_sets=self._grouping_sets, |
| ), |
| session=self._df._session, |
| ) |
| |
| agg.__doc__ = PySparkGroupedData.agg.__doc__ |
| |
| def _numeric_agg(self, function: str, cols: Sequence[str]) -> "DataFrame": |
| from pyspark.sql.connect.dataframe import DataFrame |
| |
| assert isinstance(function, str) and function in ["min", "max", "avg", "sum"] |
| |
| assert isinstance(cols, list) and all(isinstance(c, str) for c in cols) |
| |
| schema = self._df.schema |
| |
| numerical_cols: List[str] = [ |
| field.name for field in schema.fields if isinstance(field.dataType, NumericType) |
| ] |
| |
| if len(cols) > 0: |
| invalid_cols = [c for c in cols if c not in numerical_cols] |
| if len(invalid_cols) > 0: |
| raise PySparkTypeError( |
| errorClass="NOT_NUMERIC_COLUMNS", |
| messageParameters={"invalid_columns": str(invalid_cols)}, |
| ) |
| agg_cols = cols |
| else: |
| # if no column is provided, then all numerical columns are selected |
| agg_cols = numerical_cols |
| |
| return DataFrame( |
| plan.Aggregate( |
| child=self._df._plan, |
| group_type=self._group_type, |
| grouping_cols=self._grouping_cols, |
| aggregate_cols=[F._invoke_function(function, F.col(c)) for c in agg_cols], |
| pivot_col=self._pivot_col, |
| pivot_values=self._pivot_values, |
| grouping_sets=self._grouping_sets, |
| ), |
| session=self._df._session, |
| ) |
| |
| def min(self: "GroupedData", *cols: str) -> "DataFrame": |
| return self._numeric_agg("min", list(cols)) |
| |
| min.__doc__ = PySparkGroupedData.min.__doc__ |
| |
| def max(self: "GroupedData", *cols: str) -> "DataFrame": |
| return self._numeric_agg("max", list(cols)) |
| |
| max.__doc__ = PySparkGroupedData.max.__doc__ |
| |
| def sum(self: "GroupedData", *cols: str) -> "DataFrame": |
| return self._numeric_agg("sum", list(cols)) |
| |
| sum.__doc__ = PySparkGroupedData.sum.__doc__ |
| |
| def avg(self: "GroupedData", *cols: str) -> "DataFrame": |
| return self._numeric_agg("avg", list(cols)) |
| |
| avg.__doc__ = PySparkGroupedData.avg.__doc__ |
| |
| mean = avg |
| |
| def count(self: "GroupedData") -> "DataFrame": |
| return self.agg(F._invoke_function("count", F.lit(1)).alias("count")) |
| |
| count.__doc__ = PySparkGroupedData.count.__doc__ |
| |
| def pivot(self, pivot_col: str, values: Optional[List["LiteralType"]] = None) -> "GroupedData": |
| if self._group_type != "groupby": |
| if self._group_type == "pivot": |
| raise PySparkNotImplementedError( |
| errorClass="UNSUPPORTED_OPERATION", |
| messageParameters={"operation": "Repeated PIVOT operation"}, |
| ) |
| else: |
| raise PySparkNotImplementedError( |
| errorClass="UNSUPPORTED_OPERATION", |
| messageParameters={"operation": f"PIVOT after {self._group_type.upper()}"}, |
| ) |
| |
| if not isinstance(pivot_col, str): |
| raise PySparkTypeError( |
| errorClass="NOT_STR", |
| messageParameters={"arg_name": "pivot_col", "arg_type": type(pivot_col).__name__}, |
| ) |
| |
| if values is not None: |
| if not isinstance(values, list): |
| raise PySparkTypeError( |
| errorClass="NOT_LIST", |
| messageParameters={"arg_name": "values", "arg_type": type(values).__name__}, |
| ) |
| for v in values: |
| if not isinstance(v, (bool, float, int, str)): |
| raise PySparkTypeError( |
| errorClass="NOT_BOOL_OR_FLOAT_OR_INT_OR_STR", |
| messageParameters={"arg_name": "value", "arg_type": type(v).__name__}, |
| ) |
| |
| return GroupedData( |
| df=self._df, |
| group_type="pivot", |
| grouping_cols=self._grouping_cols, |
| pivot_col=self._df[pivot_col], |
| pivot_values=values, |
| ) |
| |
| pivot.__doc__ = PySparkGroupedData.pivot.__doc__ |
| |
| def apply(self, udf: "GroupedMapPandasUserDefinedFunction") -> "DataFrame": |
| # Columns are special because hasattr always return True |
| if ( |
| isinstance(udf, Column) |
| or not hasattr(udf, "func") |
| or ( |
| udf.evalType # type: ignore[attr-defined] |
| != PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF |
| ) |
| ): |
| raise PySparkTypeError( |
| errorClass="INVALID_UDF_EVAL_TYPE", |
| messageParameters={"eval_type": "SQL_GROUPED_MAP_PANDAS_UDF"}, |
| ) |
| |
| warnings.warn( |
| "It is preferred to use 'applyInPandas' over this " |
| "API. This API will be deprecated in the future releases. See SPARK-28264 for " |
| "more details.", |
| UserWarning, |
| ) |
| |
| return self.applyInPandas(udf.func, schema=udf.returnType) # type: ignore[attr-defined] |
| |
| apply.__doc__ = PySparkGroupedData.apply.__doc__ |
| |
| def applyInPandas( |
| self, func: "PandasGroupedMapFunction", schema: Union["StructType", str] |
| ) -> "DataFrame": |
| from pyspark.sql.connect.udf import UserDefinedFunction |
| from pyspark.sql.connect.dataframe import DataFrame |
| |
| _validate_vectorized_udf(func, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF) |
| if isinstance(schema, str): |
| schema = cast(StructType, self._df._session._parse_ddl(schema)) |
| udf_obj = UserDefinedFunction( |
| func, |
| returnType=schema, |
| evalType=PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF, |
| ) |
| |
| res = DataFrame( |
| plan.GroupMap( |
| child=self._df._plan, |
| grouping_cols=self._grouping_cols, |
| function=udf_obj, |
| cols=self._df.columns, |
| ), |
| session=self._df._session, |
| ) |
| if isinstance(schema, StructType): |
| res._cached_schema = schema |
| return res |
| |
| applyInPandas.__doc__ = PySparkGroupedData.applyInPandas.__doc__ |
| |
| def applyInPandasWithState( |
| self, |
| func: "PandasGroupedMapFunctionWithState", |
| outputStructType: Union[StructType, str], |
| stateStructType: Union[StructType, str], |
| outputMode: str, |
| timeoutConf: str, |
| ) -> "DataFrame": |
| from pyspark.sql.connect.udf import UserDefinedFunction |
| from pyspark.sql.connect.dataframe import DataFrame |
| |
| _validate_vectorized_udf(func, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE) |
| udf_obj = UserDefinedFunction( |
| func, |
| returnType=outputStructType, |
| evalType=PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE, |
| ) |
| |
| output_schema: str = ( |
| outputStructType.json() |
| if isinstance(outputStructType, StructType) |
| else outputStructType |
| ) |
| |
| state_schema: str = ( |
| stateStructType.json() if isinstance(stateStructType, StructType) else stateStructType |
| ) |
| |
| return DataFrame( |
| plan.ApplyInPandasWithState( |
| child=self._df._plan, |
| grouping_cols=self._grouping_cols, |
| function=udf_obj, |
| output_schema=output_schema, |
| state_schema=state_schema, |
| output_mode=outputMode, |
| timeout_conf=timeoutConf, |
| cols=self._df.columns, |
| ), |
| session=self._df._session, |
| ) |
| |
| applyInPandasWithState.__doc__ = PySparkGroupedData.applyInPandasWithState.__doc__ |
| |
| def transformWithStateInPandas( |
| self, |
| statefulProcessor: StatefulProcessor, |
| outputStructType: Union[StructType, str], |
| outputMode: str, |
| timeMode: str, |
| initialState: Optional["GroupedData"] = None, |
| eventTimeColumnName: str = "", |
| ) -> "DataFrame": |
| from pyspark.sql.connect.udf import UserDefinedFunction |
| from pyspark.sql.connect.dataframe import DataFrame |
| from pyspark.sql.streaming.stateful_processor_util import ( |
| TransformWithStateInPandasUdfUtils, |
| ) |
| |
| udf_util = TransformWithStateInPandasUdfUtils(statefulProcessor, timeMode) |
| if initialState is None: |
| udf_obj = UserDefinedFunction( |
| udf_util.transformWithStateUDF, |
| returnType=outputStructType, |
| evalType=PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_UDF, |
| ) |
| initial_state_plan = None |
| initial_state_grouping_cols = None |
| else: |
| self._df._check_same_session(initialState._df) |
| udf_obj = UserDefinedFunction( |
| udf_util.transformWithStateWithInitStateUDF, |
| returnType=outputStructType, |
| evalType=PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_INIT_STATE_UDF, |
| ) |
| initial_state_plan = initialState._df._plan |
| initial_state_grouping_cols = initialState._grouping_cols |
| |
| return DataFrame( |
| plan.TransformWithStateInPandas( |
| child=self._df._plan, |
| grouping_cols=self._grouping_cols, |
| function=udf_obj, |
| output_schema=outputStructType, |
| output_mode=outputMode, |
| time_mode=timeMode, |
| event_time_col_name=eventTimeColumnName, |
| cols=self._df.columns, |
| initial_state_plan=initial_state_plan, |
| initial_state_grouping_cols=initial_state_grouping_cols, |
| ), |
| session=self._df._session, |
| ) |
| |
| transformWithStateInPandas.__doc__ = PySparkGroupedData.transformWithStateInPandas.__doc__ |
| |
| def transformWithState( |
| self, |
| statefulProcessor: StatefulProcessor, |
| outputStructType: Union[StructType, str], |
| outputMode: str, |
| timeMode: str, |
| initialState: Optional["GroupedData"] = None, |
| eventTimeColumnName: str = "", |
| ) -> "DataFrame": |
| from pyspark.sql.connect.udf import UserDefinedFunction |
| from pyspark.sql.connect.dataframe import DataFrame |
| from pyspark.sql.streaming.stateful_processor_util import ( |
| TransformWithStateInPandasUdfUtils, |
| ) |
| |
| udf_util = TransformWithStateInPandasUdfUtils(statefulProcessor, timeMode) |
| if initialState is None: |
| udf_obj = UserDefinedFunction( |
| udf_util.transformWithStateUDF, |
| returnType=outputStructType, |
| evalType=PythonEvalType.SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_UDF, |
| ) |
| initial_state_plan = None |
| initial_state_grouping_cols = None |
| else: |
| self._df._check_same_session(initialState._df) |
| udf_obj = UserDefinedFunction( |
| udf_util.transformWithStateWithInitStateUDF, |
| returnType=outputStructType, |
| evalType=PythonEvalType.SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_INIT_STATE_UDF, |
| ) |
| initial_state_plan = initialState._df._plan |
| initial_state_grouping_cols = initialState._grouping_cols |
| |
| return DataFrame( |
| plan.TransformWithStateInPySpark( |
| child=self._df._plan, |
| grouping_cols=self._grouping_cols, |
| function=udf_obj, |
| output_schema=outputStructType, |
| output_mode=outputMode, |
| time_mode=timeMode, |
| event_time_col_name=eventTimeColumnName, |
| cols=self._df.columns, |
| initial_state_plan=initial_state_plan, |
| initial_state_grouping_cols=initial_state_grouping_cols, |
| ), |
| session=self._df._session, |
| ) |
| |
| transformWithState.__doc__ = PySparkGroupedData.transformWithState.__doc__ |
| |
| def applyInArrow( |
| self, func: "ArrowGroupedMapFunction", schema: Union[StructType, str] |
| ) -> "DataFrame": |
| from pyspark.sql.connect.udf import UserDefinedFunction |
| from pyspark.sql.connect.dataframe import DataFrame |
| |
| try: |
| # Try to infer the eval type from type hints |
| eval_type = infer_group_arrow_eval_type_from_func(func) |
| except Exception: |
| warnings.warn("Cannot infer the eval type from type hints. ", UserWarning) |
| |
| if eval_type is None: |
| eval_type = PythonEvalType.SQL_GROUPED_MAP_ARROW_UDF |
| |
| _validate_vectorized_udf(func, eval_type) |
| if isinstance(schema, str): |
| schema = cast(StructType, self._df._session._parse_ddl(schema)) |
| udf_obj = UserDefinedFunction( |
| func, |
| returnType=schema, |
| evalType=eval_type, |
| ) |
| |
| res = DataFrame( |
| plan.GroupMap( |
| child=self._df._plan, |
| grouping_cols=self._grouping_cols, |
| function=udf_obj, |
| cols=self._df.columns, |
| ), |
| session=self._df._session, |
| ) |
| if isinstance(schema, StructType): |
| res._cached_schema = schema |
| return res |
| |
| applyInArrow.__doc__ = PySparkGroupedData.applyInArrow.__doc__ |
| |
| def cogroup(self, other: "GroupedData") -> "PandasCogroupedOps": |
| return PandasCogroupedOps(self, other) |
| |
| cogroup.__doc__ = PySparkGroupedData.cogroup.__doc__ |
| |
| |
| GroupedData.__doc__ = PySparkGroupedData.__doc__ |
| |
| |
| class PandasCogroupedOps: |
| def __init__(self, gd1: "GroupedData", gd2: "GroupedData"): |
| gd1._df._check_same_session(gd2._df) |
| self._gd1 = gd1 |
| self._gd2 = gd2 |
| |
| def applyInPandas( |
| self, func: "PandasCogroupedMapFunction", schema: Union["StructType", str] |
| ) -> "DataFrame": |
| from pyspark.sql.connect.udf import UserDefinedFunction |
| from pyspark.sql.connect.dataframe import DataFrame |
| |
| _validate_vectorized_udf(func, PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF) |
| if isinstance(schema, str): |
| schema = cast(StructType, self._gd1._df._session._parse_ddl(schema)) |
| udf_obj = UserDefinedFunction( |
| func, |
| returnType=schema, |
| evalType=PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF, |
| ) |
| |
| res = DataFrame( |
| plan.CoGroupMap( |
| input=self._gd1._df._plan, |
| input_grouping_cols=self._gd1._grouping_cols, |
| other=self._gd2._df._plan, |
| other_grouping_cols=self._gd2._grouping_cols, |
| function=udf_obj, |
| ), |
| session=self._gd1._df._session, |
| ) |
| if isinstance(schema, StructType): |
| res._cached_schema = schema |
| return res |
| |
| applyInPandas.__doc__ = PySparkPandasCogroupedOps.applyInPandas.__doc__ |
| |
| def applyInArrow( |
| self, func: "ArrowCogroupedMapFunction", schema: Union[StructType, str] |
| ) -> "DataFrame": |
| from pyspark.sql.connect.udf import UserDefinedFunction |
| from pyspark.sql.connect.dataframe import DataFrame |
| |
| _validate_vectorized_udf(func, PythonEvalType.SQL_COGROUPED_MAP_ARROW_UDF) |
| if isinstance(schema, str): |
| schema = cast(StructType, self._gd1._df._session._parse_ddl(schema)) |
| udf_obj = UserDefinedFunction( |
| func, |
| returnType=schema, |
| evalType=PythonEvalType.SQL_COGROUPED_MAP_ARROW_UDF, |
| ) |
| |
| res = DataFrame( |
| plan.CoGroupMap( |
| input=self._gd1._df._plan, |
| input_grouping_cols=self._gd1._grouping_cols, |
| other=self._gd2._df._plan, |
| other_grouping_cols=self._gd2._grouping_cols, |
| function=udf_obj, |
| ), |
| session=self._gd1._df._session, |
| ) |
| if isinstance(schema, StructType): |
| res._cached_schema = schema |
| return res |
| |
| applyInArrow.__doc__ = PySparkPandasCogroupedOps.applyInArrow.__doc__ |
| |
| |
| PandasCogroupedOps.__doc__ = PySparkPandasCogroupedOps.__doc__ |
| |
| |
| def _test() -> None: |
| import os |
| import sys |
| import doctest |
| from pyspark.sql import SparkSession as PySparkSession |
| import pyspark.sql.connect.group |
| from pyspark.testing.utils import have_pandas, have_pyarrow |
| |
| globs = pyspark.sql.connect.group.__dict__.copy() |
| |
| if not have_pandas or not have_pyarrow: |
| del pyspark.sql.connect.group.GroupedData.agg.__doc__ |
| |
| globs["spark"] = ( |
| PySparkSession.builder.appName("sql.connect.group tests") |
| .remote(os.environ.get("SPARK_CONNECT_TESTING_REMOTE", "local[4]")) |
| .getOrCreate() |
| ) |
| |
| (failure_count, test_count) = doctest.testmod( |
| pyspark.sql.connect.group, |
| globs=globs, |
| optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE | doctest.REPORT_NDIFF, |
| ) |
| |
| globs["spark"].stop() |
| |
| if failure_count: |
| sys.exit(-1) |
| |
| |
| if __name__ == "__main__": |
| _test() |
