| # |
| # 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 |
| import logging |
| |
| from pyspark.sql.functions import arrow_udf, ArrowUDFType |
| from pyspark.util import PythonEvalType, is_remote_only |
| from pyspark.sql import Row |
| from pyspark.sql.types import ( |
| ArrayType, |
| YearMonthIntervalType, |
| StructType, |
| StructField, |
| VariantType, |
| VariantVal, |
| ) |
| from pyspark.sql import functions as sf |
| from pyspark.errors import AnalysisException, PythonException |
| from pyspark.testing.utils import ( |
| have_numpy, |
| numpy_requirement_message, |
| have_pyarrow, |
| pyarrow_requirement_message, |
| assertDataFrameEqual, |
| ) |
| from pyspark.testing.sqlutils import ReusedSQLTestCase |
| from typing import Iterator, Tuple |
| |
| |
| @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) |
| class GroupedAggArrowUDFTestsMixin: |
| @property |
| def data(self): |
| return ( |
| self.spark.range(10) |
| .toDF("id") |
| .withColumn("vs", sf.array([sf.lit(i * 1.0) + sf.col("id") for i in range(20, 30)])) |
| .withColumn("v", sf.explode(sf.col("vs"))) |
| .drop("vs") |
| .withColumn("w", sf.lit(1.0)) |
| ) |
| |
| @property |
| def python_plus_one(self): |
| @sf.udf("double") |
| def plus_one(v): |
| assert isinstance(v, (int, float)) |
| return float(v + 1) |
| |
| return plus_one |
| |
| @property |
| def arrow_scalar_plus_two(self): |
| import pyarrow as pa |
| |
| @arrow_udf("double", ArrowUDFType.SCALAR) |
| def plus_two(v): |
| assert isinstance(v, pa.Array) |
| return pa.compute.add(v, 2).cast(pa.float64()) |
| |
| return plus_two |
| |
| @property |
| def arrow_agg_mean_udf(self): |
| import pyarrow as pa |
| |
| @arrow_udf("double", ArrowUDFType.GROUPED_AGG) |
| def avg(v): |
| assert isinstance(v, pa.Array) |
| return pa.compute.mean(v.cast(pa.float64())) |
| |
| return avg |
| |
| @property |
| def arrow_agg_mean_arr_udf(self): |
| import pyarrow as pa |
| |
| @arrow_udf("array<double>", ArrowUDFType.GROUPED_AGG) |
| def avg(v): |
| assert isinstance(v, pa.Array) |
| assert isinstance(v, pa.ListArray) |
| return [pa.compute.mean(v.flatten()).cast(pa.float64())] |
| |
| return avg |
| |
| @property |
| def arrow_agg_sum_udf(self): |
| import pyarrow as pa |
| |
| @arrow_udf("double", ArrowUDFType.GROUPED_AGG) |
| def sum(v): |
| assert isinstance(v, pa.Array) |
| return pa.compute.sum(v).cast(pa.float64()) |
| |
| return sum |
| |
| @property |
| def arrow_agg_weighted_mean_udf(self): |
| import pyarrow as pa |
| import numpy as np |
| |
| @arrow_udf("double", ArrowUDFType.GROUPED_AGG) |
| def weighted_mean(v, w): |
| assert isinstance(v, pa.Array) |
| assert isinstance(w, pa.Array) |
| return float(np.average(v, weights=w)) |
| |
| return weighted_mean |
| |
| def test_manual(self): |
| df = self.data |
| sum_udf = self.arrow_agg_sum_udf |
| mean_udf = self.arrow_agg_mean_udf |
| mean_arr_udf = self.arrow_agg_mean_arr_udf |
| |
| result = ( |
| df.groupby("id") |
| .agg(sum_udf(df.v), mean_udf(df.v), mean_arr_udf(sf.array(df.v))) |
| .sort("id") |
| ) |
| expected = 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))"], |
| ).collect() |
| |
| self.assertEqual(expected, result.collect()) |
| |
| @unittest.skipIf(not have_numpy, numpy_requirement_message) |
| def test_basic(self): |
| df = self.data |
| weighted_mean_udf = self.arrow_agg_weighted_mean_udf |
| |
| # Groupby one column and aggregate one UDF with literal |
| result1 = df.groupby("id").agg(weighted_mean_udf(df.v, sf.lit(1.0))).sort("id") |
| expected1 = ( |
| df.groupby("id").agg(sf.mean(df.v).alias("weighted_mean(v, 1.0)")).sort("id").collect() |
| ) |
| self.assertEqual(expected1, result1.collect()) |
| |
| # Groupby one expression and aggregate one UDF with literal |
| result2 = ( |
| df.groupby((sf.col("id") + 1)).agg(weighted_mean_udf(df.v, sf.lit(1.0))).sort(df.id + 1) |
| ) |
| expected2 = ( |
| df.groupby((sf.col("id") + 1)) |
| .agg(sf.mean(df.v).alias("weighted_mean(v, 1.0)")) |
| .sort(df.id + 1) |
| ).collect() |
| self.assertEqual(expected2, result2.collect()) |
| |
| # 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(sf.mean(df.v).alias("weighted_mean(v, w)")).sort("id").collect() |
| ) |
| self.assertEqual(expected3, result3.collect()) |
| |
| # Groupby one expression and aggregate one UDF without literal |
| result4 = ( |
| df.groupby((sf.col("id") + 1).alias("id")).agg(weighted_mean_udf(df.v, df.w)).sort("id") |
| ) |
| expected4 = ( |
| df.groupby((sf.col("id") + 1).alias("id")) |
| .agg(sf.mean(df.v).alias("weighted_mean(v, w)")) |
| .sort("id") |
| ).collect() |
| self.assertEqual(expected4, result4.collect()) |
| |
| def test_alias(self): |
| df = self.data |
| mean_udf = self.arrow_agg_mean_udf |
| |
| result = df.groupby("id").agg(mean_udf(df.v).alias("mean_alias")) |
| expected = df.groupby("id").agg(sf.mean(df.v).alias("mean_alias")).collect() |
| |
| self.assertEqual(expected, result.collect()) |
| |
| def test_mixed_sql(self): |
| """ |
| Test mixing group aggregate pandas UDF with sql expression. |
| """ |
| df = self.data |
| sum_udf = self.arrow_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(sf.sum(df.v) + 1).sort("id").collect() |
| |
| # 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(sf.sum(df.v + 1)).sort("id").collect() |
| |
| # 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(sf.sum(df.v + 1) + 2).sort("id").collect() |
| |
| self.assertEqual(expected1, result1.collect()) |
| self.assertEqual(expected2, result2.collect()) |
| self.assertEqual(expected3, result3.collect()) |
| |
| 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.arrow_scalar_plus_two |
| sum_udf = self.arrow_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(sf.sum(df.v))).sort("id").collect() |
| |
| # 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(sf.sum(plus_one(df.v))).sort("id").collect() |
| |
| # 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(sf.sum(plus_two(df.v))).sort("id").collect() |
| |
| # 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(sf.sum(df.v))).sort("id").collect() |
| |
| # 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(sf.sum(plus_one(df.v)))).sort("plus_one(id)") |
| ).collect() |
| |
| # 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(sf.sum(plus_two(df.v)))).sort("plus_two(id)") |
| ).collect() |
| |
| self.assertEqual(expected1, result1.collect()) |
| self.assertEqual(expected2, result2.collect()) |
| self.assertEqual(expected3, result3.collect()) |
| self.assertEqual(expected4, result4.collect()) |
| self.assertEqual(expected5, result5.collect()) |
| self.assertEqual(expected6, result6.collect()) |
| |
| @unittest.skipIf(not have_numpy, numpy_requirement_message) |
| def test_multiple_udfs(self): |
| """ |
| Test multiple group aggregate pandas UDFs in one agg function. |
| """ |
| df = self.data |
| mean_udf = self.arrow_agg_mean_udf |
| sum_udf = self.arrow_agg_sum_udf |
| weighted_mean_udf = self.arrow_agg_weighted_mean_udf |
| |
| result = ( |
| df.groupBy("id") |
| .agg(mean_udf(df.v), sum_udf(df.v), weighted_mean_udf(df.v, df.w)) |
| .sort("id") |
| ) |
| expected = ( |
| df.groupBy("id") |
| .agg(sf.mean(df.v), sf.sum(df.v), sf.mean(df.v).alias("weighted_mean(v, w)")) |
| .sort("id") |
| .collect() |
| ) |
| |
| self.assertEqual(expected, result.collect()) |
| |
| def test_complex_groupby(self): |
| df = self.data |
| sum_udf = self.arrow_agg_sum_udf |
| plus_one = self.python_plus_one |
| plus_two = self.arrow_scalar_plus_two |
| |
| # groupby one expression |
| result1 = df.groupby(df.v % 2).agg(sum_udf(df.v)) |
| expected1 = df.groupby(df.v % 2).agg(sf.sum(df.v)).collect() |
| |
| # empty groupby |
| result2 = df.groupby().agg(sum_udf(df.v)) |
| expected2 = df.groupby().agg(sf.sum(df.v)).collect() |
| |
| # 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(sf.sum(df.v)).orderBy(df.id, df.v % 2).collect() |
| |
| # 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(sf.sum(df.v)).sort("plus_one(id)").collect() |
| |
| # 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(sf.sum(df.v)).sort("sum(v)").collect() |
| |
| # 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(sf.sum(df.v)) |
| .sort(["(v % 2)", "plus_one(id)"]) |
| ).collect() |
| |
| # 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(sf.sum(df.v)).sort(["sum(v)", "plus_two(id)"]) |
| ).collect() |
| |
| self.assertEqual(expected1, result1.collect()) |
| self.assertEqual(expected2, result2.collect()) |
| self.assertEqual(expected3, result3.collect()) |
| self.assertEqual(expected4, result4.collect()) |
| self.assertEqual(expected5, result5.collect()) |
| self.assertEqual(expected6, result6.collect()) |
| self.assertEqual(expected7, result7.collect()) |
| |
| def test_complex_expressions(self): |
| df = self.data |
| plus_one = self.python_plus_one |
| plus_two = self.arrow_scalar_plus_two |
| sum_udf = self.arrow_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(sf.col("v")), |
| sum_udf(sf.col("v1") + 3), |
| sum_udf(sf.col("v2")) + 5, |
| plus_one(sum_udf(sf.col("v1"))), |
| sum_udf(plus_one(sf.col("v2"))), |
| ) |
| .sort(["id", "(v % 2)"]) |
| ) |
| |
| expected1 = ( |
| df.withColumn("v1", df.v + 1) |
| .withColumn("v2", df.v + 2) |
| .groupby(df.id, df.v % 2) |
| .agg( |
| sf.sum(sf.col("v")), |
| sf.sum(sf.col("v1") + 3), |
| sf.sum(sf.col("v2")) + 5, |
| plus_one(sf.sum(sf.col("v1"))), |
| sf.sum(plus_one(sf.col("v2"))), |
| ) |
| .sort(["id", "(v % 2)"]) |
| .collect() |
| ) |
| |
| # 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(sf.col("v")), |
| sum_udf(sf.col("v1") + 3), |
| sum_udf(sf.col("v2")) + 5, |
| plus_two(sum_udf(sf.col("v1"))), |
| sum_udf(plus_two(sf.col("v2"))), |
| ) |
| .sort(["id", "(v % 2)"]) |
| ) |
| |
| expected2 = ( |
| df.withColumn("v1", df.v + 1) |
| .withColumn("v2", df.v + 2) |
| .groupby(df.id, df.v % 2) |
| .agg( |
| sf.sum(sf.col("v")), |
| sf.sum(sf.col("v1") + 3), |
| sf.sum(sf.col("v2")) + 5, |
| plus_two(sf.sum(sf.col("v1"))), |
| sf.sum(plus_two(sf.col("v2"))), |
| ) |
| .sort(["id", "(v % 2)"]) |
| .collect() |
| ) |
| |
| # Test sequential groupby aggregate |
| result3 = ( |
| df.groupby("id") |
| .agg(sum_udf(df.v).alias("v")) |
| .groupby("id") |
| .agg(sum_udf(sf.col("v"))) |
| .sort("id") |
| ) |
| |
| expected3 = ( |
| df.groupby("id") |
| .agg(sf.sum(df.v).alias("v")) |
| .groupby("id") |
| .agg(sf.sum(sf.col("v"))) |
| .sort("id") |
| .collect() |
| ) |
| |
| self.assertEqual(expected1, result1.collect()) |
| self.assertEqual(expected2, result2.collect()) |
| self.assertEqual(expected3, result3.collect()) |
| |
| def test_retain_group_columns(self): |
| with self.sql_conf({"spark.sql.retainGroupColumns": False}): |
| df = self.data |
| sum_udf = self.arrow_agg_sum_udf |
| |
| result1 = df.groupby(df.id).agg(sum_udf(df.v)) |
| expected1 = df.groupby(df.id).agg(sf.sum(df.v)).collect() |
| self.assertEqual(expected1, result1.collect()) |
| |
| def test_array_type(self): |
| df = self.data |
| |
| array_udf = arrow_udf(lambda x: [1.0, 2.0], "array<double>", ArrowUDFType.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): |
| with self.quiet(): |
| self.check_invalid_args() |
| |
| def check_invalid_args(self): |
| df = self.data |
| plus_one = self.python_plus_one |
| mean_udf = self.arrow_agg_mean_udf |
| with self.assertRaisesRegex(AnalysisException, "[MISSING_AGGREGATION]"): |
| df.groupby(df.id).agg(plus_one(df.v)).collect() |
| with self.assertRaisesRegex( |
| AnalysisException, "aggregate function.*argument.*aggregate function" |
| ): |
| df.groupby(df.id).agg(mean_udf(mean_udf(df.v))).collect() |
| with self.assertRaisesRegex( |
| AnalysisException, |
| "The group aggregate pandas UDF `avg` cannot be invoked together with as other, " |
| "non-pandas aggregate functions.", |
| ): |
| df.groupby(df.id).agg(mean_udf(df.v), sf.mean(df.v)).collect() |
| |
| def test_register_vectorized_udf_basic(self): |
| import pyarrow as pa |
| |
| sum_arrow_udf = arrow_udf( |
| lambda v: pa.compute.sum(v).cast(pa.int32()), |
| "integer", |
| PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF, |
| ) |
| |
| self.assertEqual(sum_arrow_udf.evalType, PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF) |
| |
| with self.temp_func("sum_arrow_udf"): |
| group_agg_pandas_udf = self.spark.udf.register("sum_arrow_udf", sum_arrow_udf) |
| self.assertEqual( |
| group_agg_pandas_udf.evalType, PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF |
| ) |
| q = """ |
| SELECT sum_arrow_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): |
| import pyarrow as pa |
| |
| 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(data).repartition(num_parts) |
| |
| f = arrow_udf(lambda x: pa.compute.sum(x).cast(pa.int32()), "int", ArrowUDFType.GROUPED_AGG) |
| |
| result = df.groupBy("id").agg(f(df["x"]).alias("sum")).sort("id") |
| self.assertEqual(result.collect(), expected) |
| |
| def test_grouped_without_group_by_clause(self): |
| import pyarrow as pa |
| |
| @arrow_udf("double", ArrowUDFType.GROUPED_AGG) |
| def max_udf(v): |
| return float(pa.compute.max(v).as_py()) |
| |
| df = self.spark.range(0, 100) |
| |
| with self.temp_view("table"), self.temp_func("max_udf"): |
| df.createTempView("table") |
| self.spark.udf.register("max_udf", max_udf) |
| |
| agg1 = df.agg(max_udf(df["id"])) |
| agg2 = self.spark.sql("select max_udf(id) from table") |
| self.assertEqual(agg1.collect(), agg2.collect()) |
| |
| def test_no_predicate_pushdown_through(self): |
| import pyarrow as pa |
| |
| @arrow_udf("float", ArrowUDFType.GROUPED_AGG) |
| def mean(x): |
| return pa.compute.mean(x).cast(pa.float32()) |
| |
| 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) |
| |
| self.assertEqual(filtered.collect()[0]["mean"], 42.0) |
| |
| assert filtered.collect()[0]["mean"] == 42.0 |
| |
| @unittest.skipIf(not have_numpy, numpy_requirement_message) |
| def test_named_arguments(self): |
| df = self.data |
| weighted_mean = self.arrow_agg_weighted_mean_udf |
| |
| with self.temp_view("v"), self.temp_func("weighted_mean"): |
| df.createOrReplaceTempView("v") |
| self.spark.udf.register("weighted_mean", weighted_mean) |
| |
| for i, aggregated in enumerate( |
| [ |
| df.groupby("id").agg(weighted_mean(df.v, w=df.w).alias("wm")), |
| df.groupby("id").agg(weighted_mean(v=df.v, w=df.w).alias("wm")), |
| df.groupby("id").agg(weighted_mean(w=df.w, v=df.v).alias("wm")), |
| self.spark.sql("SELECT id, weighted_mean(v, w => w) as wm FROM v GROUP BY id"), |
| self.spark.sql( |
| "SELECT id, weighted_mean(v => v, w => w) as wm FROM v GROUP BY id" |
| ), |
| self.spark.sql( |
| "SELECT id, weighted_mean(w => w, v => v) as wm FROM v GROUP BY id" |
| ), |
| ] |
| ): |
| with self.subTest(query_no=i): |
| self.assertEqual( |
| aggregated.collect(), |
| df.groupby("id").agg(sf.mean(df.v).alias("wm")).collect(), |
| ) |
| |
| @unittest.skipIf(not have_numpy, numpy_requirement_message) |
| def test_named_arguments_negative(self): |
| df = self.data |
| weighted_mean = self.arrow_agg_weighted_mean_udf |
| |
| with self.temp_view("v"), self.temp_func("weighted_mean"): |
| df.createOrReplaceTempView("v") |
| self.spark.udf.register("weighted_mean", weighted_mean) |
| |
| with self.assertRaisesRegex( |
| AnalysisException, |
| "DUPLICATE_ROUTINE_PARAMETER_ASSIGNMENT.DOUBLE_NAMED_ARGUMENT_REFERENCE", |
| ): |
| self.spark.sql( |
| "SELECT id, weighted_mean(v => v, v => w) as wm FROM v GROUP BY id" |
| ).show() |
| |
| with self.assertRaisesRegex(AnalysisException, "UNEXPECTED_POSITIONAL_ARGUMENT"): |
| self.spark.sql( |
| "SELECT id, weighted_mean(v => v, w) as wm FROM v GROUP BY id" |
| ).show() |
| |
| with self.assertRaisesRegex( |
| PythonException, r"weighted_mean\(\) got an unexpected keyword argument 'x'" |
| ): |
| self.spark.sql( |
| "SELECT id, weighted_mean(v => v, x => w) as wm FROM v GROUP BY id" |
| ).show() |
| |
| with self.assertRaisesRegex( |
| PythonException, r"weighted_mean\(\) got multiple values for argument 'v'" |
| ): |
| self.spark.sql( |
| "SELECT id, weighted_mean(v, v => w) as wm FROM v GROUP BY id" |
| ).show() |
| |
| def test_kwargs(self): |
| df = self.data |
| |
| @arrow_udf("double", ArrowUDFType.GROUPED_AGG) |
| def weighted_mean(**kwargs): |
| import numpy as np |
| |
| return np.average(kwargs["v"], weights=kwargs["w"]) |
| |
| with self.temp_view("v"), self.temp_func("weighted_mean"): |
| df.createOrReplaceTempView("v") |
| self.spark.udf.register("weighted_mean", weighted_mean) |
| |
| for i, aggregated in enumerate( |
| [ |
| df.groupby("id").agg(weighted_mean(v=df.v, w=df.w).alias("wm")), |
| df.groupby("id").agg(weighted_mean(w=df.w, v=df.v).alias("wm")), |
| self.spark.sql( |
| "SELECT id, weighted_mean(v => v, w => w) as wm FROM v GROUP BY id" |
| ), |
| self.spark.sql( |
| "SELECT id, weighted_mean(w => w, v => v) as wm FROM v GROUP BY id" |
| ), |
| ] |
| ): |
| with self.subTest(query_no=i): |
| self.assertEqual( |
| aggregated.collect(), |
| df.groupby("id").agg(sf.mean(df.v).alias("wm")).collect(), |
| ) |
| |
| # negative |
| with self.assertRaisesRegex( |
| AnalysisException, |
| "DUPLICATE_ROUTINE_PARAMETER_ASSIGNMENT.DOUBLE_NAMED_ARGUMENT_REFERENCE", |
| ): |
| self.spark.sql( |
| "SELECT id, weighted_mean(v => v, v => w) as wm FROM v GROUP BY id" |
| ).show() |
| |
| with self.assertRaisesRegex(AnalysisException, "UNEXPECTED_POSITIONAL_ARGUMENT"): |
| self.spark.sql( |
| "SELECT id, weighted_mean(v => v, w) as wm FROM v GROUP BY id" |
| ).show() |
| |
| def test_named_arguments_and_defaults(self): |
| import pyarrow as pa |
| |
| df = self.data |
| |
| @arrow_udf("double", ArrowUDFType.GROUPED_AGG) |
| def biased_sum(v, w=None): |
| return pa.compute.sum(v).as_py() + (pa.compute.sum(w).as_py() if w is not None else 100) |
| |
| with self.temp_view("v"), self.temp_func("biased_sum"): |
| df.createOrReplaceTempView("v") |
| self.spark.udf.register("biased_sum", biased_sum) |
| |
| # without "w" |
| for i, aggregated in enumerate( |
| [ |
| df.groupby("id").agg(biased_sum(df.v).alias("s")), |
| df.groupby("id").agg(biased_sum(v=df.v).alias("s")), |
| self.spark.sql("SELECT id, biased_sum(v) as s FROM v GROUP BY id"), |
| self.spark.sql("SELECT id, biased_sum(v => v) as s FROM v GROUP BY id"), |
| ] |
| ): |
| with self.subTest(with_w=False, query_no=i): |
| self.assertEqual( |
| aggregated.collect(), |
| df.groupby("id").agg((sf.sum(df.v) + sf.lit(100)).alias("s")).collect(), |
| ) |
| |
| # with "w" |
| for i, aggregated in enumerate( |
| [ |
| df.groupby("id").agg(biased_sum(df.v, w=df.w).alias("s")), |
| df.groupby("id").agg(biased_sum(v=df.v, w=df.w).alias("s")), |
| df.groupby("id").agg(biased_sum(w=df.w, v=df.v).alias("s")), |
| self.spark.sql("SELECT id, biased_sum(v, w => w) as s FROM v GROUP BY id"), |
| self.spark.sql("SELECT id, biased_sum(v => v, w => w) as s FROM v GROUP BY id"), |
| self.spark.sql("SELECT id, biased_sum(w => w, v => v) as s FROM v GROUP BY id"), |
| ] |
| ): |
| with self.subTest(with_w=True, query_no=i): |
| self.assertEqual( |
| aggregated.collect(), |
| df.groupby("id").agg((sf.sum(df.v) + sf.sum(df.w)).alias("s")).collect(), |
| ) |
| |
| def test_complex_agg_collect_set(self): |
| import pyarrow as pa |
| |
| df = self.spark.createDataFrame([(1, 1), (1, 2), (2, 3), (2, 5), (2, 1)], ("id", "v")) |
| |
| @arrow_udf("array<int>") |
| def arrow_collect_set(v: pa.Array) -> pa.Scalar: |
| assert isinstance(v, pa.Array), str(type(v)) |
| s = sorted([x.as_py() for x in pa.compute.unique(v)]) |
| t = pa.list_(pa.int32()) |
| return pa.scalar(value=s, type=t) |
| |
| result1 = df.select( |
| arrow_collect_set(df["id"]).alias("ids"), |
| arrow_collect_set(df["v"]).alias("vs"), |
| ) |
| |
| expected1 = df.select( |
| sf.sort_array(sf.collect_set(df["id"])).alias("ids"), |
| sf.sort_array(sf.collect_set(df["v"])).alias("vs"), |
| ) |
| |
| self.assertEqual(expected1.collect(), result1.collect()) |
| |
| def test_complex_agg_collect_list(self): |
| import pyarrow as pa |
| |
| df = self.spark.createDataFrame([(1, 1), (1, 2), (2, 3), (2, 5), (2, 1)], ("id", "v")) |
| |
| @arrow_udf("array<int>") |
| def arrow_collect_list(v: pa.Array) -> pa.Scalar: |
| assert isinstance(v, pa.Array), str(type(v)) |
| s = sorted([x.as_py() for x in v]) |
| t = pa.list_(pa.int32()) |
| return pa.scalar(value=s, type=t) |
| |
| result1 = df.select( |
| arrow_collect_list(df["id"]).alias("ids"), |
| arrow_collect_list(df["v"]).alias("vs"), |
| ) |
| |
| expected1 = df.select( |
| sf.sort_array(sf.collect_list(df["id"])).alias("ids"), |
| sf.sort_array(sf.collect_list(df["v"])).alias("vs"), |
| ) |
| |
| self.assertEqual(expected1.collect(), result1.collect()) |
| |
| def test_complex_agg_collect_as_map(self): |
| import pyarrow as pa |
| |
| df = self.spark.createDataFrame([(1, 1), (2, 2), (3, 5)], ("id", "v")) |
| |
| @arrow_udf("map<int, int>") |
| def arrow_collect_as_map(id: pa.Array, v: pa.Array) -> pa.Scalar: |
| assert isinstance(id, pa.Array), str(type(id)) |
| assert isinstance(v, pa.Array), str(type(v)) |
| d = {i: j for i, j in zip(id.to_pylist(), v.to_pylist())} |
| t = pa.map_(pa.int32(), pa.int32()) |
| return pa.scalar(value=d, type=t) |
| |
| result1 = df.select( |
| arrow_collect_as_map("id", "v").alias("map"), |
| ) |
| |
| expected1 = df.select( |
| sf.map_from_arrays(sf.collect_list("id"), sf.collect_list("v")).alias("map"), |
| ) |
| |
| self.assertEqual(expected1.collect(), result1.collect()) |
| |
| def test_complex_agg_min_max_struct(self): |
| import pyarrow as pa |
| |
| df = self.spark.createDataFrame([(1, 1), (2, 2), (3, 5)], ("id", "v")) |
| |
| @arrow_udf("struct<m1: int, m2:int>") |
| def arrow_collect_min_max(id: pa.Array, v: pa.Array) -> pa.Scalar: |
| assert isinstance(id, pa.Array), str(type(id)) |
| assert isinstance(v, pa.Array), str(type(v)) |
| m1 = pa.compute.min(id) |
| m2 = pa.compute.max(v) |
| t = pa.struct([pa.field("m1", pa.int32()), pa.field("m2", pa.int32())]) |
| return pa.scalar(value={"m1": m1.as_py(), "m2": m2.as_py()}, type=t) |
| |
| result1 = df.select( |
| arrow_collect_min_max("id", "v").alias("struct"), |
| ) |
| |
| expected1 = df.select( |
| sf.struct(sf.min("id").alias("m1"), sf.max("v").alias("m2")).alias("struct"), |
| ) |
| |
| self.assertEqual(expected1.collect(), result1.collect()) |
| |
| def test_time_min(self): |
| import pyarrow as pa |
| |
| df = self.spark.sql( |
| """ |
| SELECT * FROM VALUES |
| (1, TIME '12:34:56'), |
| (1, TIME '1:2:3'), |
| (2, TIME '0:58:59'), |
| (2, TIME '10:58:59'), |
| (2, TIME '10:00:03') |
| AS tab(i, t) |
| """ |
| ) |
| |
| @arrow_udf("time", ArrowUDFType.GROUPED_AGG) |
| def agg_min_time(v): |
| assert isinstance(v, pa.Array) |
| assert isinstance(v, pa.Time64Array) |
| return pa.compute.min(v) |
| |
| expected1 = df.select(sf.min("t").alias("res")) |
| result1 = df.select(agg_min_time("t").alias("res")) |
| self.assertEqual(expected1.collect(), result1.collect()) |
| |
| expected2 = df.groupby("i").agg(sf.min("t").alias("res")).sort("i") |
| result2 = df.groupby("i").agg(agg_min_time("t").alias("res")).sort("i") |
| self.assertEqual(expected2.collect(), result2.collect()) |
| |
| def test_input_output_variant(self): |
| import pyarrow as pa |
| |
| @arrow_udf("variant") |
| def first_variant(v: pa.Array) -> pa.Scalar: |
| assert isinstance(v, pa.Array) |
| assert isinstance(v, pa.StructArray) |
| assert isinstance(v.field("metadata"), pa.BinaryArray) |
| assert isinstance(v.field("value"), pa.BinaryArray) |
| return v[0] |
| |
| @arrow_udf("variant") |
| def last_variant(v: pa.Array) -> pa.Scalar: |
| assert isinstance(v, pa.Array) |
| assert isinstance(v, pa.StructArray) |
| assert isinstance(v.field("metadata"), pa.BinaryArray) |
| assert isinstance(v.field("value"), pa.BinaryArray) |
| return v[-1] |
| |
| df = self.spark.range(0, 10).selectExpr("parse_json(cast(id as string)) v") |
| result = df.select( |
| first_variant("v").alias("first"), |
| last_variant("v").alias("last"), |
| ) |
| self.assertEqual( |
| result.schema, |
| StructType( |
| [ |
| StructField("first", VariantType(), True), |
| StructField("last", VariantType(), True), |
| ] |
| ), |
| ) |
| |
| row = result.first() |
| self.assertIsInstance(row.first, VariantVal) |
| self.assertIsInstance(row.last, VariantVal) |
| |
| def test_return_type_coercion(self): |
| import pyarrow as pa |
| |
| df = self.spark.range(10) |
| |
| @arrow_udf("long", ArrowUDFType.GROUPED_AGG) |
| def agg_long(id: pa.Array) -> int: |
| assert isinstance(id, pa.Array), str(type(id)) |
| return pa.scalar(value=len(id), type=pa.int64()) |
| |
| result1 = df.select(agg_long("id").alias("res")) |
| self.assertEqual(1, len(result1.collect())) |
| |
| # long -> int coercion |
| @arrow_udf("int", ArrowUDFType.GROUPED_AGG) |
| def agg_int1(id: pa.Array) -> int: |
| assert isinstance(id, pa.Array), str(type(id)) |
| return pa.scalar(value=len(id), type=pa.int64()) |
| |
| result2 = df.select(agg_int1("id").alias("res")) |
| self.assertEqual(1, len(result2.collect())) |
| |
| # long -> int coercion, overflow |
| @arrow_udf("int", ArrowUDFType.GROUPED_AGG) |
| def agg_int2(id: pa.Array) -> int: |
| assert isinstance(id, pa.Array), str(type(id)) |
| return pa.scalar(value=len(id) + 2147483647, type=pa.int64()) |
| |
| result3 = df.select(agg_int2("id").alias("res")) |
| with self.assertRaises(Exception): |
| # pyarrow.lib.ArrowInvalid: |
| # Integer value 2147483657 not in range: -2147483648 to 2147483647 |
| result3.collect() |
| |
| @unittest.skipIf(not have_numpy, numpy_requirement_message) |
| def test_return_numpy_scalar(self): |
| import numpy as np |
| import pyarrow as pa |
| |
| @arrow_udf("long") |
| def np_max_udf(v: pa.Array) -> np.int64: |
| assert isinstance(v, pa.Array) |
| return np.max(v) |
| |
| @arrow_udf("long") |
| def np_min_udf(v: pa.Array) -> np.int64: |
| assert isinstance(v, pa.Array) |
| return np.min(v) |
| |
| @arrow_udf("double") |
| def np_avg_udf(v: pa.Array) -> np.float64: |
| assert isinstance(v, pa.Array) |
| return np.mean(v) |
| |
| df = self.spark.range(10) |
| expected = df.select( |
| sf.max("id").alias("max"), |
| sf.min("id").alias("min"), |
| sf.avg("id").alias("avg"), |
| ) |
| |
| result = df.select( |
| np_max_udf("id").alias("max"), |
| np_min_udf("id").alias("min"), |
| np_avg_udf("id").alias("avg"), |
| ) |
| self.assertEqual(expected.collect(), result.collect()) |
| |
| def test_unsupported_return_types(self): |
| import pyarrow as pa |
| |
| with self.quiet(): |
| with self.assertRaisesRegex( |
| NotImplementedError, |
| "Invalid return type with grouped aggregate " |
| "Arrow UDFs.*ArrayType.*YearMonthIntervalType", |
| ): |
| arrow_udf( |
| lambda x: x, |
| ArrayType(ArrayType(YearMonthIntervalType())), |
| ArrowUDFType.GROUPED_AGG, |
| ) |
| |
| with self.assertRaisesRegex( |
| NotImplementedError, |
| "Invalid return type with grouped aggregate " |
| "Arrow UDFs.*ArrayType.*YearMonthIntervalType", |
| ): |
| |
| @arrow_udf(ArrayType(ArrayType(YearMonthIntervalType())), ArrowUDFType.GROUPED_AGG) |
| def func_a(a: pa.Array) -> pa.Scalar: |
| return pa.compute.max(a) |
| |
| def test_0_args(self): |
| import pyarrow as pa |
| |
| df = self.spark.range(10).withColumn("k", sf.col("id") % 3) |
| |
| @arrow_udf("long", ArrowUDFType.GROUPED_AGG) |
| def arrow_max(v) -> int: |
| return pa.compute.max(v).as_py() |
| |
| @arrow_udf("long", ArrowUDFType.GROUPED_AGG) |
| def arrow_lit_1() -> int: |
| return 1 |
| |
| expected1 = df.select(sf.max("id").alias("res1"), sf.lit(1).alias("res1")) |
| result1 = df.select(arrow_max("id").alias("res1"), arrow_lit_1().alias("res1")) |
| self.assertEqual(expected1.collect(), result1.collect()) |
| |
| expected2 = ( |
| df.groupby("k") |
| .agg( |
| sf.max("id").alias("res1"), |
| sf.lit(1).alias("res1"), |
| ) |
| .sort("k") |
| ) |
| result2 = ( |
| df.groupby("k") |
| .agg( |
| arrow_max("id").alias("res1"), |
| arrow_lit_1().alias("res1"), |
| ) |
| .sort("k") |
| ) |
| self.assertEqual(expected2.collect(), result2.collect()) |
| |
| def test_arrow_batch_slicing(self): |
| import pyarrow as pa |
| |
| df = self.spark.range(10000000).select( |
| (sf.col("id") % 2).alias("key"), sf.col("id").alias("v") |
| ) |
| |
| @arrow_udf("long", ArrowUDFType.GROUPED_AGG) |
| def arrow_max(v): |
| assert len(v) == 10000000 / 2, len(v) |
| return pa.compute.max(v) |
| |
| expected = (df.groupby("key").agg(sf.max("v").alias("res")).sort("key")).collect() |
| |
| for maxRecords, maxBytes in [(1000, 2**31 - 1), (0, 1048576), (1000, 1048576)]: |
| with self.subTest(maxRecords=maxRecords, maxBytes=maxBytes): |
| with self.sql_conf( |
| { |
| "spark.sql.execution.arrow.maxRecordsPerBatch": maxRecords, |
| "spark.sql.execution.arrow.maxBytesPerBatch": maxBytes, |
| } |
| ): |
| result = ( |
| df.groupBy("key").agg(arrow_max("v").alias("res")).sort("key") |
| ).collect() |
| |
| self.assertEqual(expected, result) |
| |
| @unittest.skipIf(is_remote_only(), "Requires JVM access") |
| def test_grouped_agg_arrow_udf_with_logging(self): |
| import pyarrow as pa |
| |
| @arrow_udf("double", ArrowUDFType.GROUPED_AGG) |
| def my_grouped_agg_arrow_udf(x): |
| assert isinstance(x, pa.Array) |
| logger = logging.getLogger("test_grouped_agg_arrow") |
| logger.warning(f"grouped agg arrow udf: {len(x)}") |
| return pa.compute.sum(x) |
| |
| df = self.spark.createDataFrame( |
| [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], ("id", "v") |
| ) |
| |
| with self.sql_conf({"spark.sql.pyspark.worker.logging.enabled": "true"}): |
| assertDataFrameEqual( |
| df.groupby("id").agg(my_grouped_agg_arrow_udf("v").alias("result")), |
| [Row(id=1, result=3.0), Row(id=2, result=18.0)], |
| ) |
| |
| logs = self.spark.tvf.python_worker_logs() |
| |
| assertDataFrameEqual( |
| logs.select("level", "msg", "context", "logger"), |
| [ |
| Row( |
| level="WARNING", |
| msg=f"grouped agg arrow udf: {n}", |
| context={"func_name": my_grouped_agg_arrow_udf.__name__}, |
| logger="test_grouped_agg_arrow", |
| ) |
| for n in [2, 3] |
| ], |
| ) |
| |
| def test_iterator_grouped_agg_single_column(self): |
| """ |
| Test iterator API for grouped aggregation with single column. |
| """ |
| import pyarrow as pa |
| from typing import Iterator |
| |
| @arrow_udf("double") |
| def arrow_mean_iter(it: Iterator[pa.Array]) -> float: |
| sum_val = 0.0 |
| cnt = 0 |
| for v in it: |
| assert isinstance(v, pa.Array) |
| sum_val += pa.compute.sum(v).as_py() |
| cnt += len(v) |
| return sum_val / cnt if cnt > 0 else 0.0 |
| |
| df = self.spark.createDataFrame( |
| [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], ("id", "v") |
| ) |
| |
| result = df.groupby("id").agg(arrow_mean_iter(df["v"]).alias("mean")).sort("id") |
| expected = df.groupby("id").agg(sf.mean(df["v"]).alias("mean")).sort("id").collect() |
| |
| self.assertEqual(expected, result.collect()) |
| |
| @unittest.skipIf(not have_numpy, numpy_requirement_message) |
| def test_iterator_grouped_agg_multiple_columns(self): |
| """ |
| Test iterator API for grouped aggregation with multiple columns. |
| """ |
| import pyarrow as pa |
| import numpy as np |
| |
| @arrow_udf("double") |
| def arrow_weighted_mean_iter(it: Iterator[Tuple[pa.Array, pa.Array]]) -> float: |
| weighted_sum = 0.0 |
| weight = 0.0 |
| for v, w in it: |
| assert isinstance(v, pa.Array) |
| assert isinstance(w, pa.Array) |
| weighted_sum += np.dot(v, w) |
| weight += pa.compute.sum(w).as_py() |
| return weighted_sum / weight if weight > 0 else 0.0 |
| |
| df = self.spark.createDataFrame( |
| [(1, 1.0, 1.0), (1, 2.0, 2.0), (2, 3.0, 1.0), (2, 5.0, 2.0), (2, 10.0, 3.0)], |
| ("id", "v", "w"), |
| ) |
| |
| result = ( |
| df.groupby("id") |
| .agg(arrow_weighted_mean_iter(df["v"], df["w"]).alias("wm")) |
| .sort("id") |
| .collect() |
| ) |
| |
| # Expected weighted means: |
| # Group 1: (1.0*1.0 + 2.0*2.0) / (1.0 + 2.0) = 5.0 / 3.0 |
| # Group 2: (3.0*1.0 + 5.0*2.0 + 10.0*3.0) / (1.0 + 2.0 + 3.0) = 43.0 / 6.0 |
| expected = [(1, 5.0 / 3.0), (2, 43.0 / 6.0)] |
| |
| self.assertEqual(len(result), len(expected)) |
| for r, (exp_id, exp_wm) in zip(result, expected): |
| self.assertEqual(r["id"], exp_id) |
| self.assertAlmostEqual(r["wm"], exp_wm, places=5) |
| |
| def test_iterator_grouped_agg_eval_type(self): |
| """ |
| Test that the eval type is correctly inferred for iterator grouped agg UDFs. |
| """ |
| import pyarrow as pa |
| from typing import Iterator |
| |
| @arrow_udf("double") |
| def arrow_sum_iter(it: Iterator[pa.Array]) -> float: |
| total = 0.0 |
| for v in it: |
| total += pa.compute.sum(v).as_py() |
| return total |
| |
| self.assertEqual(arrow_sum_iter.evalType, PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF) |
| |
| def test_iterator_grouped_agg_partial_consumption(self): |
| """ |
| Test that iterator grouped agg UDF can partially consume batches. |
| This ensures that batches are processed one by one without loading all data into memory. |
| """ |
| import pyarrow as pa |
| from typing import Iterator |
| |
| # Create a dataset with multiple batches per group |
| # Use small batch size to ensure multiple batches per group |
| # Use same value for all data points to avoid ordering issues |
| with self.sql_conf({"spark.sql.execution.arrow.maxRecordsPerBatch": 2}): |
| df = self.spark.createDataFrame( |
| [(1, 1.0), (1, 1.0), (1, 1.0), (1, 1.0), (2, 1.0), (2, 1.0)], ("id", "v") |
| ) |
| |
| @arrow_udf("struct<count:bigint,sum:double>") |
| def arrow_count_sum_partial(it: Iterator[pa.Array]) -> dict: |
| # Only consume first two batches, then return |
| # This tests that partial consumption works correctly |
| total = 0.0 |
| count = 0 |
| for i, v in enumerate(it): |
| if i < 2: # Only process first 2 batches |
| total += pa.compute.sum(v).as_py() |
| count += len(v) |
| else: |
| # Stop early - partial consumption |
| break |
| return {"count": count, "sum": total} |
| |
| result = ( |
| df.groupby("id").agg(arrow_count_sum_partial(df["v"]).alias("result")).sort("id") |
| ) |
| |
| # Verify results are correct for partial consumption |
| # With batch size = 2: |
| # Group 1 (id=1): 4 values in 2 batches -> processes both batches |
| # Batch 1: [1.0, 1.0], Batch 2: [1.0, 1.0] |
| # Result: count=4, sum=4.0 |
| # Group 2 (id=2): 2 values in 1 batch -> processes 1 batch (only 1 batch available) |
| # Batch 1: [1.0, 1.0] |
| # Result: count=2, sum=2.0 |
| actual = result.collect() |
| self.assertEqual(len(actual), 2, "Should have results for both groups") |
| |
| # Verify both groups were processed correctly |
| # Group 1: processes 2 batches (all available) |
| group1_result = next(row for row in actual if row["id"] == 1) |
| self.assertEqual( |
| group1_result["result"]["count"], |
| 4, |
| msg="Group 1 should process 4 values (2 batches)", |
| ) |
| self.assertAlmostEqual( |
| group1_result["result"]["sum"], 4.0, places=5, msg="Group 1 should sum to 4.0" |
| ) |
| |
| # Group 2: processes 1 batch (only batch available) |
| group2_result = next(row for row in actual if row["id"] == 2) |
| self.assertEqual( |
| group2_result["result"]["count"], |
| 2, |
| msg="Group 2 should process 2 values (1 batch)", |
| ) |
| self.assertAlmostEqual( |
| group2_result["result"]["sum"], 2.0, places=5, msg="Group 2 should sum to 2.0" |
| ) |
| |
| def test_iterator_grouped_agg_sql_single_column(self): |
| """ |
| Test iterator API for grouped aggregation with single column in SQL. |
| """ |
| import pyarrow as pa |
| |
| @arrow_udf("double") |
| def arrow_mean_iter(it: Iterator[pa.Array]) -> float: |
| sum_val = 0.0 |
| cnt = 0 |
| for v in it: |
| assert isinstance(v, pa.Array) |
| sum_val += pa.compute.sum(v).as_py() |
| cnt += len(v) |
| return sum_val / cnt if cnt > 0 else 0.0 |
| |
| df = self.spark.createDataFrame( |
| [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], ("id", "v") |
| ) |
| |
| with self.temp_view("test_table"), self.temp_func("arrow_mean_iter"): |
| df.createOrReplaceTempView("test_table") |
| self.spark.udf.register("arrow_mean_iter", arrow_mean_iter) |
| |
| # Test SQL query with GROUP BY |
| result_sql = self.spark.sql( |
| "SELECT id, arrow_mean_iter(v) as mean FROM test_table GROUP BY id ORDER BY id" |
| ) |
| expected = df.groupby("id").agg(sf.mean(df["v"]).alias("mean")).sort("id").collect() |
| |
| self.assertEqual(expected, result_sql.collect()) |
| |
| def test_iterator_grouped_agg_sql_multiple_columns(self): |
| """ |
| Test iterator API for grouped aggregation with multiple columns in SQL. |
| """ |
| import pyarrow as pa |
| |
| @arrow_udf("double") |
| def arrow_weighted_mean_iter(it: Iterator[Tuple[pa.Array, pa.Array]]) -> float: |
| weighted_sum = 0.0 |
| weight = 0.0 |
| for v, w in it: |
| assert isinstance(v, pa.Array) |
| assert isinstance(w, pa.Array) |
| weighted_sum += pa.compute.sum(pa.compute.multiply(v, w)).as_py() |
| weight += pa.compute.sum(w).as_py() |
| return weighted_sum / weight if weight > 0 else 0.0 |
| |
| df = self.spark.createDataFrame( |
| [(1, 1.0, 1.0), (1, 2.0, 2.0), (2, 3.0, 1.0), (2, 5.0, 2.0), (2, 10.0, 3.0)], |
| ("id", "v", "w"), |
| ) |
| |
| with self.temp_view("test_table"), self.temp_func("arrow_weighted_mean_iter"): |
| df.createOrReplaceTempView("test_table") |
| self.spark.udf.register("arrow_weighted_mean_iter", arrow_weighted_mean_iter) |
| |
| # Test SQL query with GROUP BY and multiple columns |
| result_sql = self.spark.sql( |
| "SELECT id, arrow_weighted_mean_iter(v, w) as wm " |
| "FROM test_table GROUP BY id ORDER BY id" |
| ) |
| |
| # Expected weighted means: |
| # Group 1: (1.0*1.0 + 2.0*2.0) / (1.0 + 2.0) = 5.0 / 3.0 |
| # Group 2: (3.0*1.0 + 5.0*2.0 + 10.0*3.0) / (1.0 + 2.0 + 3.0) = 43.0 / 6.0 |
| expected = [Row(id=1, wm=5.0 / 3.0), Row(id=2, wm=43.0 / 6.0)] |
| |
| actual_results = result_sql.collect() |
| self.assertEqual(actual_results, expected) |
| |
| |
| class GroupedAggArrowUDFTests(GroupedAggArrowUDFTestsMixin, ReusedSQLTestCase): |
| pass |
| |
| |
| if __name__ == "__main__": |
| from pyspark.testing import main |
| |
| main() |
