Coverage for pyspark/sql/tests/test_pandas_udf_grouped_agg.py : 20%

# 

# 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. 

# 

import unittest 

from pyspark.rdd import PythonEvalType 

from pyspark.sql import Row 

from pyspark.sql.functions import array, explode, col, lit, mean, sum, \ 

udf, pandas_udf, PandasUDFType 

from pyspark.sql.types import ArrayType, TimestampType 

from pyspark.sql.utils import AnalysisException 

from pyspark.testing.sqlutils import ReusedSQLTestCase, have_pandas, have_pyarrow, \ 

pandas_requirement_message, pyarrow_requirement_message 

from pyspark.testing.utils import QuietTest 

if have_pandas: 

import pandas as pd 

from pandas.testing import assert_frame_equal 

@unittest.skipIf( 

not have_pandas or not have_pyarrow, 

pandas_requirement_message or pyarrow_requirement_message) # type: ignore[arg-type] 

class GroupedAggPandasUDFTests(ReusedSQLTestCase): 

@property 

def data(self): 

return self.spark.range(10).toDF('id') \ 

.withColumn("vs", array([lit(i * 1.0) + col('id') for i in range(20, 30)])) \ 

.withColumn("v", explode(col('vs'))) \ 

.drop('vs') \ 

.withColumn('w', lit(1.0)) 

@property 

def python_plus_one(self): 

@udf('double') 

def plus_one(v): 

assert isinstance(v, (int, float)) 

return float(v + 1) 

return plus_one 

@property 

def pandas_scalar_plus_two(self): 

@pandas_udf('double', PandasUDFType.SCALAR) 

def plus_two(v): 

assert isinstance(v, pd.Series) 

return v + 2 

return plus_two 

@property 

def pandas_agg_mean_udf(self): 

@pandas_udf('double', PandasUDFType.GROUPED_AGG) 

def avg(v): 

return v.mean() 

return avg 

@property 

def pandas_agg_sum_udf(self): 

@pandas_udf('double', PandasUDFType.GROUPED_AGG) 

def sum(v): 

return v.sum() 

return sum 

@property 

def pandas_agg_weighted_mean_udf(self): 

import numpy as np 

@pandas_udf('double', PandasUDFType.GROUPED_AGG) 

def weighted_mean(v, w): 

return np.average(v, weights=w) 

return weighted_mean 

def test_manual(self): 

df = self.data 

sum_udf = self.pandas_agg_sum_udf 

mean_udf = self.pandas_agg_mean_udf 

mean_arr_udf = pandas_udf( 

self.pandas_agg_mean_udf.func, 

ArrayType(self.pandas_agg_mean_udf.returnType), 

self.pandas_agg_mean_udf.evalType) 

result1 = df.groupby('id').agg( 

sum_udf(df.v), 

mean_udf(df.v), 

mean_arr_udf(array(df.v))).sort('id') 

expected1 = self.spark.createDataFrame( 

[[0, 245.0, 24.5, [24.5]], 

[1, 255.0, 25.5, [25.5]], 

[2, 265.0, 26.5, [26.5]], 

[3, 275.0, 27.5, [27.5]], 

[4, 285.0, 28.5, [28.5]], 

[5, 295.0, 29.5, [29.5]], 

[6, 305.0, 30.5, [30.5]], 

[7, 315.0, 31.5, [31.5]], 

[8, 325.0, 32.5, [32.5]], 

[9, 335.0, 33.5, [33.5]]], 

['id', 'sum(v)', 'avg(v)', 'avg(array(v))']) 

assert_frame_equal(expected1.toPandas(), result1.toPandas()) 

def test_basic(self): 

df = self.data 

weighted_mean_udf = self.pandas_agg_weighted_mean_udf 

# Groupby one column and aggregate one UDF with literal 

result1 = df.groupby('id').agg(weighted_mean_udf(df.v, lit(1.0))).sort('id') 

expected1 = df.groupby('id').agg(mean(df.v).alias('weighted_mean(v, 1.0)')).sort('id') 

assert_frame_equal(expected1.toPandas(), result1.toPandas()) 

# Groupby one expression and aggregate one UDF with literal 

result2 = df.groupby((col('id') + 1)).agg(weighted_mean_udf(df.v, lit(1.0)))\ 

.sort(df.id + 1) 

expected2 = df.groupby((col('id') + 1))\ 

.agg(mean(df.v).alias('weighted_mean(v, 1.0)')).sort(df.id + 1) 

assert_frame_equal(expected2.toPandas(), result2.toPandas()) 

# Groupby one column and aggregate one UDF without literal 

result3 = df.groupby('id').agg(weighted_mean_udf(df.v, df.w)).sort('id') 

expected3 = df.groupby('id').agg(mean(df.v).alias('weighted_mean(v, w)')).sort('id') 

assert_frame_equal(expected3.toPandas(), result3.toPandas()) 

# Groupby one expression and aggregate one UDF without literal 

result4 = df.groupby((col('id') + 1).alias('id'))\ 

.agg(weighted_mean_udf(df.v, df.w))\ 

.sort('id') 

expected4 = df.groupby((col('id') + 1).alias('id'))\ 

.agg(mean(df.v).alias('weighted_mean(v, w)'))\ 

.sort('id') 

assert_frame_equal(expected4.toPandas(), result4.toPandas()) 

def test_unsupported_types(self): 

with QuietTest(self.sc): 

with self.assertRaisesRegex(NotImplementedError, 'not supported'): 

pandas_udf( 

lambda x: x, 

ArrayType(ArrayType(TimestampType())), 

PandasUDFType.GROUPED_AGG) 

with QuietTest(self.sc): 

with self.assertRaisesRegex(NotImplementedError, 'not supported'): 

@pandas_udf('mean double, std double', PandasUDFType.GROUPED_AGG) 

def mean_and_std_udf(v): 

return v.mean(), v.std() 

with QuietTest(self.sc): 

with self.assertRaisesRegex(NotImplementedError, 'not supported'): 

@pandas_udf(ArrayType(TimestampType()), PandasUDFType.GROUPED_AGG) 

def mean_and_std_udf(v): 

return {v.mean(): v.std()} 

def test_alias(self): 

df = self.data 

mean_udf = self.pandas_agg_mean_udf 

result1 = df.groupby('id').agg(mean_udf(df.v).alias('mean_alias')) 

expected1 = df.groupby('id').agg(mean(df.v).alias('mean_alias')) 

assert_frame_equal(expected1.toPandas(), result1.toPandas()) 

def test_mixed_sql(self): 

""" 

Test mixing group aggregate pandas UDF with sql expression. 

""" 

df = self.data 

sum_udf = self.pandas_agg_sum_udf 

# Mix group aggregate pandas UDF with sql expression 

result1 = (df.groupby('id') 

.agg(sum_udf(df.v) + 1) 

.sort('id')) 

expected1 = (df.groupby('id') 

.agg(sum(df.v) + 1) 

.sort('id')) 

# Mix group aggregate pandas UDF with sql expression (order swapped) 

result2 = (df.groupby('id') 

.agg(sum_udf(df.v + 1)) 

.sort('id')) 

expected2 = (df.groupby('id') 

.agg(sum(df.v + 1)) 

.sort('id')) 

# Wrap group aggregate pandas UDF with two sql expressions 

result3 = (df.groupby('id') 

.agg(sum_udf(df.v + 1) + 2) 

.sort('id')) 

expected3 = (df.groupby('id') 

.agg(sum(df.v + 1) + 2) 

.sort('id')) 

assert_frame_equal(expected1.toPandas(), result1.toPandas()) 

assert_frame_equal(expected2.toPandas(), result2.toPandas()) 

assert_frame_equal(expected3.toPandas(), result3.toPandas()) 

def test_mixed_udfs(self): 

""" 

Test mixing group aggregate pandas UDF with python UDF and scalar pandas UDF. 

""" 

df = self.data 

plus_one = self.python_plus_one 

plus_two = self.pandas_scalar_plus_two 

sum_udf = self.pandas_agg_sum_udf 

# Mix group aggregate pandas UDF and python UDF 

result1 = (df.groupby('id') 

.agg(plus_one(sum_udf(df.v))) 

.sort('id')) 

expected1 = (df.groupby('id') 

.agg(plus_one(sum(df.v))) 

.sort('id')) 

# Mix group aggregate pandas UDF and python UDF (order swapped) 

result2 = (df.groupby('id') 

.agg(sum_udf(plus_one(df.v))) 

.sort('id')) 

expected2 = (df.groupby('id') 

.agg(sum(plus_one(df.v))) 

.sort('id')) 

# Mix group aggregate pandas UDF and scalar pandas UDF 

result3 = (df.groupby('id') 

.agg(sum_udf(plus_two(df.v))) 

.sort('id')) 

expected3 = (df.groupby('id') 

.agg(sum(plus_two(df.v))) 

.sort('id')) 

# Mix group aggregate pandas UDF and scalar pandas UDF (order swapped) 

result4 = (df.groupby('id') 

.agg(plus_two(sum_udf(df.v))) 

.sort('id')) 

expected4 = (df.groupby('id') 

.agg(plus_two(sum(df.v))) 

.sort('id')) 

# Wrap group aggregate pandas UDF with two python UDFs and use python UDF in groupby 

result5 = (df.groupby(plus_one(df.id)) 

.agg(plus_one(sum_udf(plus_one(df.v)))) 

.sort('plus_one(id)')) 

expected5 = (df.groupby(plus_one(df.id)) 

.agg(plus_one(sum(plus_one(df.v)))) 

.sort('plus_one(id)')) 

# Wrap group aggregate pandas UDF with two scala pandas UDF and user scala pandas UDF in 

# groupby 

result6 = (df.groupby(plus_two(df.id)) 

.agg(plus_two(sum_udf(plus_two(df.v)))) 

.sort('plus_two(id)')) 

expected6 = (df.groupby(plus_two(df.id)) 

.agg(plus_two(sum(plus_two(df.v)))) 

.sort('plus_two(id)')) 

assert_frame_equal(expected1.toPandas(), result1.toPandas()) 

assert_frame_equal(expected2.toPandas(), result2.toPandas()) 

assert_frame_equal(expected3.toPandas(), result3.toPandas()) 

assert_frame_equal(expected4.toPandas(), result4.toPandas()) 

assert_frame_equal(expected5.toPandas(), result5.toPandas()) 

assert_frame_equal(expected6.toPandas(), result6.toPandas()) 

def test_multiple_udfs(self): 

""" 

Test multiple group aggregate pandas UDFs in one agg function. 

""" 

df = self.data 

mean_udf = self.pandas_agg_mean_udf 

sum_udf = self.pandas_agg_sum_udf 

weighted_mean_udf = self.pandas_agg_weighted_mean_udf 

result1 = (df.groupBy('id') 

.agg(mean_udf(df.v), 

sum_udf(df.v), 

weighted_mean_udf(df.v, df.w)) 

.sort('id') 

.toPandas()) 

expected1 = (df.groupBy('id') 

.agg(mean(df.v), 

sum(df.v), 

mean(df.v).alias('weighted_mean(v, w)')) 

.sort('id') 

.toPandas()) 

assert_frame_equal(expected1, result1) 

def test_complex_groupby(self): 

df = self.data 

sum_udf = self.pandas_agg_sum_udf 

plus_one = self.python_plus_one 

plus_two = self.pandas_scalar_plus_two 

# groupby one expression 

result1 = df.groupby(df.v % 2).agg(sum_udf(df.v)) 

expected1 = df.groupby(df.v % 2).agg(sum(df.v)) 

# empty groupby 

result2 = df.groupby().agg(sum_udf(df.v)) 

expected2 = df.groupby().agg(sum(df.v)) 

# groupby one column and one sql expression 

result3 = df.groupby(df.id, df.v % 2).agg(sum_udf(df.v)).orderBy(df.id, df.v % 2) 

expected3 = df.groupby(df.id, df.v % 2).agg(sum(df.v)).orderBy(df.id, df.v % 2) 

# groupby one python UDF 

result4 = df.groupby(plus_one(df.id)).agg(sum_udf(df.v)).sort('plus_one(id)') 

expected4 = df.groupby(plus_one(df.id)).agg(sum(df.v)).sort('plus_one(id)') 

# groupby one scalar pandas UDF 

result5 = df.groupby(plus_two(df.id)).agg(sum_udf(df.v)).sort('sum(v)') 

expected5 = df.groupby(plus_two(df.id)).agg(sum(df.v)).sort('sum(v)') 

# groupby one expression and one python UDF 

result6 = (df.groupby(df.v % 2, plus_one(df.id)) 

.agg(sum_udf(df.v)).sort(['(v % 2)', 'plus_one(id)'])) 

expected6 = (df.groupby(df.v % 2, plus_one(df.id)) 

.agg(sum(df.v)).sort(['(v % 2)', 'plus_one(id)'])) 

# groupby one expression and one scalar pandas UDF 

result7 = (df.groupby(df.v % 2, plus_two(df.id)) 

.agg(sum_udf(df.v)).sort(['sum(v)', 'plus_two(id)'])) 

expected7 = (df.groupby(df.v % 2, plus_two(df.id)) 

.agg(sum(df.v)).sort(['sum(v)', 'plus_two(id)'])) 

assert_frame_equal(expected1.toPandas(), result1.toPandas()) 

assert_frame_equal(expected2.toPandas(), result2.toPandas()) 

assert_frame_equal(expected3.toPandas(), result3.toPandas()) 

assert_frame_equal(expected4.toPandas(), result4.toPandas()) 

assert_frame_equal(expected5.toPandas(), result5.toPandas()) 

assert_frame_equal(expected6.toPandas(), result6.toPandas()) 

assert_frame_equal(expected7.toPandas(), result7.toPandas()) 

def test_complex_expressions(self): 

df = self.data 

plus_one = self.python_plus_one 

plus_two = self.pandas_scalar_plus_two 

sum_udf = self.pandas_agg_sum_udf 

# Test complex expressions with sql expression, python UDF and 

# group aggregate pandas UDF 

result1 = (df.withColumn('v1', plus_one(df.v)) 

.withColumn('v2', df.v + 2) 

.groupby(df.id, df.v % 2) 

.agg(sum_udf(col('v')), 

sum_udf(col('v1') + 3), 

sum_udf(col('v2')) + 5, 

plus_one(sum_udf(col('v1'))), 

sum_udf(plus_one(col('v2')))) 

.sort(['id', '(v % 2)']) 

.toPandas().sort_values(by=['id', '(v % 2)'])) 

expected1 = (df.withColumn('v1', df.v + 1) 

.withColumn('v2', df.v + 2) 

.groupby(df.id, df.v % 2) 

.agg(sum(col('v')), 

sum(col('v1') + 3), 

sum(col('v2')) + 5, 

plus_one(sum(col('v1'))), 

sum(plus_one(col('v2')))) 

.sort(['id', '(v % 2)']) 

.toPandas().sort_values(by=['id', '(v % 2)'])) 

# Test complex expressions with sql expression, scala pandas UDF and 

# group aggregate pandas UDF 

result2 = (df.withColumn('v1', plus_one(df.v)) 

.withColumn('v2', df.v + 2) 

.groupby(df.id, df.v % 2) 

.agg(sum_udf(col('v')), 

sum_udf(col('v1') + 3), 

sum_udf(col('v2')) + 5, 

plus_two(sum_udf(col('v1'))), 

sum_udf(plus_two(col('v2')))) 

.sort(['id', '(v % 2)']) 

.toPandas().sort_values(by=['id', '(v % 2)'])) 

expected2 = (df.withColumn('v1', df.v + 1) 

.withColumn('v2', df.v + 2) 

.groupby(df.id, df.v % 2) 

.agg(sum(col('v')), 

sum(col('v1') + 3), 

sum(col('v2')) + 5, 

plus_two(sum(col('v1'))), 

sum(plus_two(col('v2')))) 

.sort(['id', '(v % 2)']) 

.toPandas().sort_values(by=['id', '(v % 2)'])) 

# Test sequential groupby aggregate 

result3 = (df.groupby('id') 

.agg(sum_udf(df.v).alias('v')) 

.groupby('id') 

.agg(sum_udf(col('v'))) 

.sort('id') 

.toPandas()) 

expected3 = (df.groupby('id') 

.agg(sum(df.v).alias('v')) 

.groupby('id') 

.agg(sum(col('v'))) 

.sort('id') 

.toPandas()) 

assert_frame_equal(expected1, result1) 

assert_frame_equal(expected2, result2) 

assert_frame_equal(expected3, result3) 

def test_retain_group_columns(self): 

with self.sql_conf({"spark.sql.retainGroupColumns": False}): 

df = self.data 

sum_udf = self.pandas_agg_sum_udf 

result1 = df.groupby(df.id).agg(sum_udf(df.v)) 

expected1 = df.groupby(df.id).agg(sum(df.v)) 

assert_frame_equal(expected1.toPandas(), result1.toPandas()) 

def test_array_type(self): 

df = self.data 

array_udf = pandas_udf(lambda x: [1.0, 2.0], 'array<double>', PandasUDFType.GROUPED_AGG) 

result1 = df.groupby('id').agg(array_udf(df['v']).alias('v2')) 

self.assertEqual(result1.first()['v2'], [1.0, 2.0]) 

def test_invalid_args(self): 

df = self.data 

plus_one = self.python_plus_one 

mean_udf = self.pandas_agg_mean_udf 

with QuietTest(self.sc): 

with self.assertRaisesRegex( 

AnalysisException, 

'nor.*aggregate function'): 

df.groupby(df.id).agg(plus_one(df.v)).collect() 

with QuietTest(self.sc): 

with self.assertRaisesRegex( 

AnalysisException, 

'aggregate function.*argument.*aggregate function'): 

df.groupby(df.id).agg(mean_udf(mean_udf(df.v))).collect() 

with QuietTest(self.sc): 

with self.assertRaisesRegex( 

AnalysisException, 

'mixture.*aggregate function.*group aggregate pandas UDF'): 

df.groupby(df.id).agg(mean_udf(df.v), mean(df.v)).collect() 

def test_register_vectorized_udf_basic(self): 

sum_pandas_udf = pandas_udf( 

lambda v: v.sum(), "integer", PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF) 

self.assertEqual(sum_pandas_udf.evalType, PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF) 

group_agg_pandas_udf = self.spark.udf.register("sum_pandas_udf", sum_pandas_udf) 

self.assertEqual(group_agg_pandas_udf.evalType, PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF) 

q = "SELECT sum_pandas_udf(v1) FROM VALUES (3, 0), (2, 0), (1, 1) tbl(v1, v2) GROUP BY v2" 

actual = sorted(map(lambda r: r[0], self.spark.sql(q).collect())) 

expected = [1, 5] 

self.assertEqual(actual, expected) 

def test_grouped_with_empty_partition(self): 

data = [Row(id=1, x=2), Row(id=1, x=3), Row(id=2, x=4)] 

expected = [Row(id=1, sum=5), Row(id=2, x=4)] 

num_parts = len(data) + 1 

df = self.spark.createDataFrame(self.sc.parallelize(data, numSlices=num_parts)) 

f = pandas_udf(lambda x: x.sum(), 

'int', PandasUDFType.GROUPED_AGG) 

result = df.groupBy('id').agg(f(df['x']).alias('sum')).collect() 

self.assertEqual(result, expected) 

def test_grouped_without_group_by_clause(self): 

@pandas_udf('double', PandasUDFType.GROUPED_AGG) 

def max_udf(v): 

return v.max() 

df = self.spark.range(0, 100) 

self.spark.udf.register('max_udf', max_udf) 

with self.tempView("table"): 

df.createTempView('table') 

agg1 = df.agg(max_udf(df['id'])) 

agg2 = self.spark.sql("select max_udf(id) from table") 

assert_frame_equal(agg1.toPandas(), agg2.toPandas()) 

def test_no_predicate_pushdown_through(self): 

# SPARK-30921: We should not pushdown predicates of PythonUDFs through Aggregate. 

import numpy as np 

@pandas_udf('float', PandasUDFType.GROUPED_AGG) 

def mean(x): 

return np.mean(x) 

df = self.spark.createDataFrame([ 

Row(id=1, foo=42), Row(id=2, foo=1), Row(id=2, foo=2) 

]) 

agg = df.groupBy('id').agg(mean('foo').alias("mean")) 

filtered = agg.filter(agg['mean'] > 40.0) 

assert(filtered.collect()[0]["mean"] == 42.0) 

if __name__ == "__main__": 

from pyspark.sql.tests.test_pandas_udf_grouped_agg import * # noqa: F401 

try: 

import xmlrunner # type: ignore[import] 

testRunner = xmlrunner.XMLTestRunner(output='target/test-reports', verbosity=2) 

except ImportError: 

testRunner = None 

unittest.main(testRunner=testRunner, verbosity=2)