| # |
| # 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 datetime |
| from itertools import chain |
| import re |
| import math |
| |
| from py4j.protocol import Py4JJavaError |
| from pyspark.sql import Row, Window, types |
| from pyspark.sql.functions import ( |
| udf, |
| input_file_name, |
| col, |
| percentile_approx, |
| lit, |
| assert_true, |
| sum_distinct, |
| sumDistinct, |
| shiftleft, |
| shiftLeft, |
| shiftRight, |
| shiftright, |
| shiftrightunsigned, |
| shiftRightUnsigned, |
| octet_length, |
| bit_length, |
| sec, |
| csc, |
| cot, |
| make_date, |
| date_add, |
| date_sub, |
| add_months, |
| array_repeat, |
| size, |
| slice, |
| least, |
| ) |
| from pyspark.testing.sqlutils import ReusedSQLTestCase, SQLTestUtils |
| |
| |
| class FunctionsTests(ReusedSQLTestCase): |
| def test_explode(self): |
| from pyspark.sql.functions import explode, explode_outer, posexplode_outer |
| |
| d = [ |
| Row(a=1, intlist=[1, 2, 3], mapfield={"a": "b"}), |
| Row(a=1, intlist=[], mapfield={}), |
| Row(a=1, intlist=None, mapfield=None), |
| ] |
| rdd = self.sc.parallelize(d) |
| data = self.spark.createDataFrame(rdd) |
| |
| result = data.select(explode(data.intlist).alias("a")).select("a").collect() |
| self.assertEqual(result[0][0], 1) |
| self.assertEqual(result[1][0], 2) |
| self.assertEqual(result[2][0], 3) |
| |
| result = data.select(explode(data.mapfield).alias("a", "b")).select("a", "b").collect() |
| self.assertEqual(result[0][0], "a") |
| self.assertEqual(result[0][1], "b") |
| |
| result = [tuple(x) for x in data.select(posexplode_outer("intlist")).collect()] |
| self.assertEqual(result, [(0, 1), (1, 2), (2, 3), (None, None), (None, None)]) |
| |
| result = [tuple(x) for x in data.select(posexplode_outer("mapfield")).collect()] |
| self.assertEqual(result, [(0, "a", "b"), (None, None, None), (None, None, None)]) |
| |
| result = [x[0] for x in data.select(explode_outer("intlist")).collect()] |
| self.assertEqual(result, [1, 2, 3, None, None]) |
| |
| result = [tuple(x) for x in data.select(explode_outer("mapfield")).collect()] |
| self.assertEqual(result, [("a", "b"), (None, None), (None, None)]) |
| |
| def test_basic_functions(self): |
| rdd = self.sc.parallelize(['{"foo":"bar"}', '{"foo":"baz"}']) |
| df = self.spark.read.json(rdd) |
| df.count() |
| df.collect() |
| df.schema |
| |
| # cache and checkpoint |
| self.assertFalse(df.is_cached) |
| df.persist() |
| df.unpersist(True) |
| df.cache() |
| self.assertTrue(df.is_cached) |
| self.assertEqual(2, df.count()) |
| |
| with self.tempView("temp"): |
| df.createOrReplaceTempView("temp") |
| df = self.spark.sql("select foo from temp") |
| df.count() |
| df.collect() |
| |
| def test_corr(self): |
| import math |
| |
| df = self.sc.parallelize([Row(a=i, b=math.sqrt(i)) for i in range(10)]).toDF() |
| corr = df.stat.corr("a", "b") |
| self.assertTrue(abs(corr - 0.95734012) < 1e-6) |
| |
| def test_sampleby(self): |
| df = self.sc.parallelize([Row(a=i, b=(i % 3)) for i in range(100)]).toDF() |
| sampled = df.stat.sampleBy("b", fractions={0: 0.5, 1: 0.5}, seed=0) |
| self.assertTrue(sampled.count() == 35) |
| |
| def test_cov(self): |
| df = self.sc.parallelize([Row(a=i, b=2 * i) for i in range(10)]).toDF() |
| cov = df.stat.cov("a", "b") |
| self.assertTrue(abs(cov - 55.0 / 3) < 1e-6) |
| |
| def test_crosstab(self): |
| df = self.sc.parallelize([Row(a=i % 3, b=i % 2) for i in range(1, 7)]).toDF() |
| ct = df.stat.crosstab("a", "b").collect() |
| ct = sorted(ct, key=lambda x: x[0]) |
| for i, row in enumerate(ct): |
| self.assertEqual(row[0], str(i)) |
| self.assertTrue(row[1], 1) |
| self.assertTrue(row[2], 1) |
| |
| def test_math_functions(self): |
| df = self.sc.parallelize([Row(a=i, b=2 * i) for i in range(10)]).toDF() |
| from pyspark.sql import functions |
| |
| SQLTestUtils.assert_close( |
| [math.cos(i) for i in range(10)], df.select(functions.cos(df.a)).collect() |
| ) |
| SQLTestUtils.assert_close( |
| [math.cos(i) for i in range(10)], df.select(functions.cos("a")).collect() |
| ) |
| SQLTestUtils.assert_close( |
| [math.sin(i) for i in range(10)], df.select(functions.sin(df.a)).collect() |
| ) |
| SQLTestUtils.assert_close( |
| [math.sin(i) for i in range(10)], df.select(functions.sin(df["a"])).collect() |
| ) |
| SQLTestUtils.assert_close( |
| [math.pow(i, 2 * i) for i in range(10)], df.select(functions.pow(df.a, df.b)).collect() |
| ) |
| SQLTestUtils.assert_close( |
| [math.pow(i, 2) for i in range(10)], df.select(functions.pow(df.a, 2)).collect() |
| ) |
| SQLTestUtils.assert_close( |
| [math.pow(i, 2) for i in range(10)], df.select(functions.pow(df.a, 2.0)).collect() |
| ) |
| SQLTestUtils.assert_close( |
| [math.hypot(i, 2 * i) for i in range(10)], |
| df.select(functions.hypot(df.a, df.b)).collect(), |
| ) |
| SQLTestUtils.assert_close( |
| [math.hypot(i, 2 * i) for i in range(10)], |
| df.select(functions.hypot("a", "b")).collect(), |
| ) |
| SQLTestUtils.assert_close( |
| [math.hypot(i, 2) for i in range(10)], df.select(functions.hypot("a", 2)).collect() |
| ) |
| SQLTestUtils.assert_close( |
| [math.hypot(i, 2) for i in range(10)], df.select(functions.hypot(df.a, 2)).collect() |
| ) |
| |
| def test_inverse_trig_functions(self): |
| from pyspark.sql import functions |
| |
| funs = [ |
| (functions.acosh, "ACOSH"), |
| (functions.asinh, "ASINH"), |
| (functions.atanh, "ATANH"), |
| ] |
| |
| cols = ["a", functions.col("a")] |
| |
| for f, alias in funs: |
| for c in cols: |
| self.assertIn(f"{alias}(a)", repr(f(c))) |
| |
| def test_reciprocal_trig_functions(self): |
| # SPARK-36683: Tests for reciprocal trig functions (SEC, CSC and COT) |
| lst = [ |
| 0.0, |
| math.pi / 6, |
| math.pi / 4, |
| math.pi / 3, |
| math.pi / 2, |
| math.pi, |
| 3 * math.pi / 2, |
| 2 * math.pi, |
| ] |
| |
| df = self.spark.createDataFrame(lst, types.DoubleType()) |
| |
| def to_reciprocal_trig(func): |
| return [1.0 / func(i) if func(i) != 0 else math.inf for i in lst] |
| |
| SQLTestUtils.assert_close(to_reciprocal_trig(math.cos), df.select(sec(df.value)).collect()) |
| SQLTestUtils.assert_close(to_reciprocal_trig(math.sin), df.select(csc(df.value)).collect()) |
| SQLTestUtils.assert_close(to_reciprocal_trig(math.tan), df.select(cot(df.value)).collect()) |
| |
| def test_rand_functions(self): |
| df = self.df |
| from pyspark.sql import functions |
| |
| rnd = df.select("key", functions.rand()).collect() |
| for row in rnd: |
| assert row[1] >= 0.0 and row[1] <= 1.0, "got: %s" % row[1] |
| rndn = df.select("key", functions.randn(5)).collect() |
| for row in rndn: |
| assert row[1] >= -4.0 and row[1] <= 4.0, "got: %s" % row[1] |
| |
| # If the specified seed is 0, we should use it. |
| # https://issues.apache.org/jira/browse/SPARK-9691 |
| rnd1 = df.select("key", functions.rand(0)).collect() |
| rnd2 = df.select("key", functions.rand(0)).collect() |
| self.assertEqual(sorted(rnd1), sorted(rnd2)) |
| |
| rndn1 = df.select("key", functions.randn(0)).collect() |
| rndn2 = df.select("key", functions.randn(0)).collect() |
| self.assertEqual(sorted(rndn1), sorted(rndn2)) |
| |
| def test_string_functions(self): |
| from pyspark.sql import functions |
| from pyspark.sql.functions import col, lit |
| |
| string_functions = [ |
| "upper", |
| "lower", |
| "ascii", |
| "base64", |
| "unbase64", |
| "ltrim", |
| "rtrim", |
| "trim", |
| ] |
| |
| df = self.spark.createDataFrame([["nick"]], schema=["name"]) |
| self.assertRaisesRegex( |
| TypeError, "must be the same type", lambda: df.select(col("name").substr(0, lit(1))) |
| ) |
| |
| for name in string_functions: |
| self.assertEqual( |
| df.select(getattr(functions, name)("name")).first()[0], |
| df.select(getattr(functions, name)(col("name"))).first()[0], |
| ) |
| |
| def test_octet_length_function(self): |
| # SPARK-36751: add octet length api for python |
| df = self.spark.createDataFrame([("cat",), ("\U0001F408",)], ["cat"]) |
| actual = df.select(octet_length("cat")).collect() |
| self.assertEqual([Row(3), Row(4)], actual) |
| |
| def test_bit_length_function(self): |
| # SPARK-36751: add bit length api for python |
| df = self.spark.createDataFrame([("cat",), ("\U0001F408",)], ["cat"]) |
| actual = df.select(bit_length("cat")).collect() |
| self.assertEqual([Row(24), Row(32)], actual) |
| |
| def test_array_contains_function(self): |
| from pyspark.sql.functions import array_contains |
| |
| df = self.spark.createDataFrame([(["1", "2", "3"],), ([],)], ["data"]) |
| actual = df.select(array_contains(df.data, "1").alias("b")).collect() |
| self.assertEqual([Row(b=True), Row(b=False)], actual) |
| |
| def test_between_function(self): |
| df = self.sc.parallelize( |
| [Row(a=1, b=2, c=3), Row(a=2, b=1, c=3), Row(a=4, b=1, c=4)] |
| ).toDF() |
| self.assertEqual( |
| [Row(a=2, b=1, c=3), Row(a=4, b=1, c=4)], df.filter(df.a.between(df.b, df.c)).collect() |
| ) |
| |
| def test_dayofweek(self): |
| from pyspark.sql.functions import dayofweek |
| |
| dt = datetime.datetime(2017, 11, 6) |
| df = self.spark.createDataFrame([Row(date=dt)]) |
| row = df.select(dayofweek(df.date)).first() |
| self.assertEqual(row[0], 2) |
| |
| # Test added for SPARK-37738; change Python API to accept both col & int as input |
| def test_date_add_function(self): |
| dt = datetime.date(2021, 12, 27) |
| |
| # Note; number var in Python gets converted to LongType column; |
| # this is not supported by the function, so cast to Integer explicitly |
| df = self.spark.createDataFrame([Row(date=dt, add=2)], "date date, add integer") |
| |
| self.assertTrue( |
| all( |
| df.select( |
| date_add(df.date, df.add) == datetime.date(2021, 12, 29), |
| date_add(df.date, "add") == datetime.date(2021, 12, 29), |
| date_add(df.date, 3) == datetime.date(2021, 12, 30), |
| ).first() |
| ) |
| ) |
| |
| # Test added for SPARK-37738; change Python API to accept both col & int as input |
| def test_date_sub_function(self): |
| dt = datetime.date(2021, 12, 27) |
| |
| # Note; number var in Python gets converted to LongType column; |
| # this is not supported by the function, so cast to Integer explicitly |
| df = self.spark.createDataFrame([Row(date=dt, sub=2)], "date date, sub integer") |
| |
| self.assertTrue( |
| all( |
| df.select( |
| date_sub(df.date, df.sub) == datetime.date(2021, 12, 25), |
| date_sub(df.date, "sub") == datetime.date(2021, 12, 25), |
| date_sub(df.date, 3) == datetime.date(2021, 12, 24), |
| ).first() |
| ) |
| ) |
| |
| # Test added for SPARK-37738; change Python API to accept both col & int as input |
| def test_add_months_function(self): |
| dt = datetime.date(2021, 12, 27) |
| |
| # Note; number in Python gets converted to LongType column; |
| # this is not supported by the function, so cast to Integer explicitly |
| df = self.spark.createDataFrame([Row(date=dt, add=2)], "date date, add integer") |
| |
| self.assertTrue( |
| all( |
| df.select( |
| add_months(df.date, df.add) == datetime.date(2022, 2, 27), |
| add_months(df.date, "add") == datetime.date(2022, 2, 27), |
| add_months(df.date, 3) == datetime.date(2022, 3, 27), |
| ).first() |
| ) |
| ) |
| |
| def test_make_date(self): |
| # SPARK-36554: expose make_date expression |
| df = self.spark.createDataFrame([(2020, 6, 26)], ["Y", "M", "D"]) |
| row_from_col = df.select(make_date(df.Y, df.M, df.D)).first() |
| self.assertEqual(row_from_col[0], datetime.date(2020, 6, 26)) |
| row_from_name = df.select(make_date("Y", "M", "D")).first() |
| self.assertEqual(row_from_name[0], datetime.date(2020, 6, 26)) |
| |
| def test_expr(self): |
| from pyspark.sql import functions |
| |
| row = Row(a="length string", b=75) |
| df = self.spark.createDataFrame([row]) |
| result = df.select(functions.expr("length(a)")).collect()[0].asDict() |
| self.assertEqual(13, result["length(a)"]) |
| |
| # add test for SPARK-10577 (test broadcast join hint) |
| def test_functions_broadcast(self): |
| from pyspark.sql.functions import broadcast |
| |
| df1 = self.spark.createDataFrame([(1, "1"), (2, "2")], ("key", "value")) |
| df2 = self.spark.createDataFrame([(1, "1"), (2, "2")], ("key", "value")) |
| |
| # equijoin - should be converted into broadcast join |
| plan1 = df1.join(broadcast(df2), "key")._jdf.queryExecution().executedPlan() |
| self.assertEqual(1, plan1.toString().count("BroadcastHashJoin")) |
| |
| # no join key -- should not be a broadcast join |
| plan2 = df1.crossJoin(broadcast(df2))._jdf.queryExecution().executedPlan() |
| self.assertEqual(0, plan2.toString().count("BroadcastHashJoin")) |
| |
| # planner should not crash without a join |
| broadcast(df1)._jdf.queryExecution().executedPlan() |
| |
| def test_first_last_ignorenulls(self): |
| from pyspark.sql import functions |
| |
| df = self.spark.range(0, 100) |
| df2 = df.select(functions.when(df.id % 3 == 0, None).otherwise(df.id).alias("id")) |
| df3 = df2.select( |
| functions.first(df2.id, False).alias("a"), |
| functions.first(df2.id, True).alias("b"), |
| functions.last(df2.id, False).alias("c"), |
| functions.last(df2.id, True).alias("d"), |
| ) |
| self.assertEqual([Row(a=None, b=1, c=None, d=98)], df3.collect()) |
| |
| def test_approxQuantile(self): |
| df = self.sc.parallelize([Row(a=i, b=i + 10) for i in range(10)]).toDF() |
| for f in ["a", "a"]: |
| aq = df.stat.approxQuantile(f, [0.1, 0.5, 0.9], 0.1) |
| self.assertTrue(isinstance(aq, list)) |
| self.assertEqual(len(aq), 3) |
| self.assertTrue(all(isinstance(q, float) for q in aq)) |
| aqs = df.stat.approxQuantile(["a", "b"], [0.1, 0.5, 0.9], 0.1) |
| self.assertTrue(isinstance(aqs, list)) |
| self.assertEqual(len(aqs), 2) |
| self.assertTrue(isinstance(aqs[0], list)) |
| self.assertEqual(len(aqs[0]), 3) |
| self.assertTrue(all(isinstance(q, float) for q in aqs[0])) |
| self.assertTrue(isinstance(aqs[1], list)) |
| self.assertEqual(len(aqs[1]), 3) |
| self.assertTrue(all(isinstance(q, float) for q in aqs[1])) |
| aqt = df.stat.approxQuantile(("a", "b"), [0.1, 0.5, 0.9], 0.1) |
| self.assertTrue(isinstance(aqt, list)) |
| self.assertEqual(len(aqt), 2) |
| self.assertTrue(isinstance(aqt[0], list)) |
| self.assertEqual(len(aqt[0]), 3) |
| self.assertTrue(all(isinstance(q, float) for q in aqt[0])) |
| self.assertTrue(isinstance(aqt[1], list)) |
| self.assertEqual(len(aqt[1]), 3) |
| self.assertTrue(all(isinstance(q, float) for q in aqt[1])) |
| self.assertRaises(TypeError, lambda: df.stat.approxQuantile(123, [0.1, 0.9], 0.1)) |
| self.assertRaises(TypeError, lambda: df.stat.approxQuantile(("a", 123), [0.1, 0.9], 0.1)) |
| self.assertRaises(TypeError, lambda: df.stat.approxQuantile(["a", 123], [0.1, 0.9], 0.1)) |
| |
| def test_sorting_functions_with_column(self): |
| from pyspark.sql import functions |
| from pyspark.sql.column import Column |
| |
| funs = [ |
| functions.asc_nulls_first, |
| functions.asc_nulls_last, |
| functions.desc_nulls_first, |
| functions.desc_nulls_last, |
| ] |
| exprs = [col("x"), "x"] |
| |
| for fun in funs: |
| for expr in exprs: |
| res = fun(expr) |
| self.assertIsInstance(res, Column) |
| self.assertIn(f"""'x {fun.__name__.replace("_", " ").upper()}'""", str(res)) |
| |
| for expr in exprs: |
| res = functions.asc(expr) |
| self.assertIsInstance(res, Column) |
| self.assertIn("""'x ASC NULLS FIRST'""", str(res)) |
| |
| for expr in exprs: |
| res = functions.desc(expr) |
| self.assertIsInstance(res, Column) |
| self.assertIn("""'x DESC NULLS LAST'""", str(res)) |
| |
| def test_sort_with_nulls_order(self): |
| from pyspark.sql import functions |
| |
| df = self.spark.createDataFrame( |
| [("Tom", 80), (None, 60), ("Alice", 50)], ["name", "height"] |
| ) |
| self.assertEqual( |
| df.select(df.name).orderBy(functions.asc_nulls_first("name")).collect(), |
| [Row(name=None), Row(name="Alice"), Row(name="Tom")], |
| ) |
| self.assertEqual( |
| df.select(df.name).orderBy(functions.asc_nulls_last("name")).collect(), |
| [Row(name="Alice"), Row(name="Tom"), Row(name=None)], |
| ) |
| self.assertEqual( |
| df.select(df.name).orderBy(functions.desc_nulls_first("name")).collect(), |
| [Row(name=None), Row(name="Tom"), Row(name="Alice")], |
| ) |
| self.assertEqual( |
| df.select(df.name).orderBy(functions.desc_nulls_last("name")).collect(), |
| [Row(name="Tom"), Row(name="Alice"), Row(name=None)], |
| ) |
| |
| def test_input_file_name_reset_for_rdd(self): |
| rdd = self.sc.textFile("python/test_support/hello/hello.txt").map(lambda x: {"data": x}) |
| df = self.spark.createDataFrame(rdd, "data STRING") |
| df.select(input_file_name().alias("file")).collect() |
| |
| non_file_df = self.spark.range(100).select(input_file_name()) |
| |
| results = non_file_df.collect() |
| self.assertTrue(len(results) == 100) |
| |
| # [SPARK-24605]: if everything was properly reset after the last job, this should return |
| # empty string rather than the file read in the last job. |
| for result in results: |
| self.assertEqual(result[0], "") |
| |
| def test_slice(self): |
| df = self.spark.createDataFrame( |
| [ |
| ( |
| [1, 2, 3], |
| 2, |
| 2, |
| ), |
| ( |
| [4, 5], |
| 2, |
| 2, |
| ), |
| ], |
| ["x", "index", "len"], |
| ) |
| |
| expected = [Row(sliced=[2, 3]), Row(sliced=[5])] |
| self.assertTrue( |
| all( |
| [ |
| df.select(slice(df.x, 2, 2).alias("sliced")).collect() == expected, |
| df.select(slice(df.x, lit(2), lit(2)).alias("sliced")).collect() == expected, |
| df.select(slice("x", "index", "len").alias("sliced")).collect() == expected, |
| ] |
| ) |
| ) |
| |
| self.assertEqual( |
| df.select(slice(df.x, size(df.x) - 1, lit(1)).alias("sliced")).collect(), |
| [Row(sliced=[2]), Row(sliced=[4])], |
| ) |
| self.assertEqual( |
| df.select(slice(df.x, lit(1), size(df.x) - 1).alias("sliced")).collect(), |
| [Row(sliced=[1, 2]), Row(sliced=[4])], |
| ) |
| |
| def test_array_repeat(self): |
| df = self.spark.range(1) |
| df = df.withColumn("repeat_n", lit(3)) |
| |
| expected = [Row(val=[0, 0, 0])] |
| self.assertTrue( |
| all( |
| [ |
| df.select(array_repeat("id", 3).alias("val")).collect() == expected, |
| df.select(array_repeat("id", lit(3)).alias("val")).collect() == expected, |
| df.select(array_repeat("id", "repeat_n").alias("val")).collect() == expected, |
| ] |
| ) |
| ) |
| |
| def test_input_file_name_udf(self): |
| df = self.spark.read.text("python/test_support/hello/hello.txt") |
| df = df.select(udf(lambda x: x)("value"), input_file_name().alias("file")) |
| file_name = df.collect()[0].file |
| self.assertTrue("python/test_support/hello/hello.txt" in file_name) |
| |
| def test_least(self): |
| df = self.spark.createDataFrame([(1, 4, 3)], ["a", "b", "c"]) |
| |
| expected = [Row(least=1)] |
| self.assertTrue( |
| all( |
| [ |
| df.select(least(df.a, df.b, df.c).alias("least")).collect() == expected, |
| df.select(least(lit(3), lit(5), lit(1)).alias("least")).collect() == expected, |
| df.select(least("a", "b", "c").alias("least")).collect() == expected, |
| ] |
| ) |
| ) |
| |
| def test_overlay(self): |
| from pyspark.sql.functions import col, lit, overlay |
| from itertools import chain |
| import re |
| |
| actual = list( |
| chain.from_iterable( |
| [ |
| re.findall("(overlay\\(.*\\))", str(x)) |
| for x in [ |
| overlay(col("foo"), col("bar"), 1), |
| overlay("x", "y", 3), |
| overlay(col("x"), col("y"), 1, 3), |
| overlay("x", "y", 2, 5), |
| overlay("x", "y", lit(11)), |
| overlay("x", "y", lit(2), lit(5)), |
| ] |
| ] |
| ) |
| ) |
| |
| expected = [ |
| "overlay(foo, bar, 1, -1)", |
| "overlay(x, y, 3, -1)", |
| "overlay(x, y, 1, 3)", |
| "overlay(x, y, 2, 5)", |
| "overlay(x, y, 11, -1)", |
| "overlay(x, y, 2, 5)", |
| ] |
| |
| self.assertListEqual(actual, expected) |
| |
| df = self.spark.createDataFrame([("SPARK_SQL", "CORE", 7, 0)], ("x", "y", "pos", "len")) |
| |
| exp = [Row(ol="SPARK_CORESQL")] |
| self.assertTrue( |
| all( |
| [ |
| df.select(overlay(df.x, df.y, 7, 0).alias("ol")).collect() == exp, |
| df.select(overlay(df.x, df.y, lit(7), lit(0)).alias("ol")).collect() == exp, |
| df.select(overlay("x", "y", "pos", "len").alias("ol")).collect() == exp, |
| ] |
| ) |
| ) |
| |
| def test_percentile_approx(self): |
| actual = list( |
| chain.from_iterable( |
| [ |
| re.findall("(percentile_approx\\(.*\\))", str(x)) |
| for x in [ |
| percentile_approx(col("foo"), lit(0.5)), |
| percentile_approx(col("bar"), 0.25, 42), |
| percentile_approx(col("bar"), [0.25, 0.5, 0.75]), |
| percentile_approx(col("foo"), (0.05, 0.95), 100), |
| percentile_approx("foo", 0.5), |
| percentile_approx("bar", [0.1, 0.9], lit(10)), |
| ] |
| ] |
| ) |
| ) |
| |
| expected = [ |
| "percentile_approx(foo, 0.5, 10000)", |
| "percentile_approx(bar, 0.25, 42)", |
| "percentile_approx(bar, array(0.25, 0.5, 0.75), 10000)", |
| "percentile_approx(foo, array(0.05, 0.95), 100)", |
| "percentile_approx(foo, 0.5, 10000)", |
| "percentile_approx(bar, array(0.1, 0.9), 10)", |
| ] |
| |
| self.assertListEqual(actual, expected) |
| |
| def test_nth_value(self): |
| from pyspark.sql import Window |
| from pyspark.sql.functions import nth_value |
| |
| df = self.spark.createDataFrame( |
| [ |
| ("a", 0, None), |
| ("a", 1, "x"), |
| ("a", 2, "y"), |
| ("a", 3, "z"), |
| ("a", 4, None), |
| ("b", 1, None), |
| ("b", 2, None), |
| ], |
| schema=("key", "order", "value"), |
| ) |
| w = Window.partitionBy("key").orderBy("order") |
| |
| rs = df.select( |
| df.key, |
| df.order, |
| nth_value("value", 2).over(w), |
| nth_value("value", 2, False).over(w), |
| nth_value("value", 2, True).over(w), |
| ).collect() |
| |
| expected = [ |
| ("a", 0, None, None, None), |
| ("a", 1, "x", "x", None), |
| ("a", 2, "x", "x", "y"), |
| ("a", 3, "x", "x", "y"), |
| ("a", 4, "x", "x", "y"), |
| ("b", 1, None, None, None), |
| ("b", 2, None, None, None), |
| ] |
| |
| for r, ex in zip(sorted(rs), sorted(expected)): |
| self.assertEqual(tuple(r), ex[: len(r)]) |
| |
| def test_higher_order_function_failures(self): |
| from pyspark.sql.functions import col, transform |
| |
| # Should fail with varargs |
| with self.assertRaises(ValueError): |
| transform(col("foo"), lambda *x: lit(1)) |
| |
| # Should fail with kwargs |
| with self.assertRaises(ValueError): |
| transform(col("foo"), lambda **x: lit(1)) |
| |
| # Should fail with nullary function |
| with self.assertRaises(ValueError): |
| transform(col("foo"), lambda: lit(1)) |
| |
| # Should fail with quaternary function |
| with self.assertRaises(ValueError): |
| transform(col("foo"), lambda x1, x2, x3, x4: lit(1)) |
| |
| # Should fail if function doesn't return Column |
| with self.assertRaises(ValueError): |
| transform(col("foo"), lambda x: 1) |
| |
| def test_nested_higher_order_function(self): |
| # SPARK-35382: lambda vars must be resolved properly in nested higher order functions |
| from pyspark.sql.functions import flatten, struct, transform |
| |
| df = self.spark.sql("SELECT array(1, 2, 3) as numbers, array('a', 'b', 'c') as letters") |
| |
| actual = df.select( |
| flatten( |
| transform( |
| "numbers", |
| lambda number: transform( |
| "letters", lambda letter: struct(number.alias("n"), letter.alias("l")) |
| ), |
| ) |
| ) |
| ).first()[0] |
| |
| expected = [ |
| (1, "a"), |
| (1, "b"), |
| (1, "c"), |
| (2, "a"), |
| (2, "b"), |
| (2, "c"), |
| (3, "a"), |
| (3, "b"), |
| (3, "c"), |
| ] |
| |
| self.assertEquals(actual, expected) |
| |
| def test_window_functions(self): |
| df = self.spark.createDataFrame([(1, "1"), (2, "2"), (1, "2"), (1, "2")], ["key", "value"]) |
| w = Window.partitionBy("value").orderBy("key") |
| from pyspark.sql import functions as F |
| |
| sel = df.select( |
| df.value, |
| df.key, |
| F.max("key").over(w.rowsBetween(0, 1)), |
| F.min("key").over(w.rowsBetween(0, 1)), |
| F.count("key").over(w.rowsBetween(float("-inf"), float("inf"))), |
| F.row_number().over(w), |
| F.rank().over(w), |
| F.dense_rank().over(w), |
| F.ntile(2).over(w), |
| ) |
| rs = sorted(sel.collect()) |
| expected = [ |
| ("1", 1, 1, 1, 1, 1, 1, 1, 1), |
| ("2", 1, 1, 1, 3, 1, 1, 1, 1), |
| ("2", 1, 2, 1, 3, 2, 1, 1, 1), |
| ("2", 2, 2, 2, 3, 3, 3, 2, 2), |
| ] |
| for r, ex in zip(rs, expected): |
| self.assertEqual(tuple(r), ex[: len(r)]) |
| |
| def test_window_functions_without_partitionBy(self): |
| df = self.spark.createDataFrame([(1, "1"), (2, "2"), (1, "2"), (1, "2")], ["key", "value"]) |
| w = Window.orderBy("key", df.value) |
| from pyspark.sql import functions as F |
| |
| sel = df.select( |
| df.value, |
| df.key, |
| F.max("key").over(w.rowsBetween(0, 1)), |
| F.min("key").over(w.rowsBetween(0, 1)), |
| F.count("key").over(w.rowsBetween(float("-inf"), float("inf"))), |
| F.row_number().over(w), |
| F.rank().over(w), |
| F.dense_rank().over(w), |
| F.ntile(2).over(w), |
| ) |
| rs = sorted(sel.collect()) |
| expected = [ |
| ("1", 1, 1, 1, 4, 1, 1, 1, 1), |
| ("2", 1, 1, 1, 4, 2, 2, 2, 1), |
| ("2", 1, 2, 1, 4, 3, 2, 2, 2), |
| ("2", 2, 2, 2, 4, 4, 4, 3, 2), |
| ] |
| for r, ex in zip(rs, expected): |
| self.assertEqual(tuple(r), ex[: len(r)]) |
| |
| def test_window_functions_cumulative_sum(self): |
| df = self.spark.createDataFrame([("one", 1), ("two", 2)], ["key", "value"]) |
| from pyspark.sql import functions as F |
| |
| # Test cumulative sum |
| sel = df.select( |
| df.key, F.sum(df.value).over(Window.rowsBetween(Window.unboundedPreceding, 0)) |
| ) |
| rs = sorted(sel.collect()) |
| expected = [("one", 1), ("two", 3)] |
| for r, ex in zip(rs, expected): |
| self.assertEqual(tuple(r), ex[: len(r)]) |
| |
| # Test boundary values less than JVM's Long.MinValue and make sure we don't overflow |
| sel = df.select( |
| df.key, F.sum(df.value).over(Window.rowsBetween(Window.unboundedPreceding - 1, 0)) |
| ) |
| rs = sorted(sel.collect()) |
| expected = [("one", 1), ("two", 3)] |
| for r, ex in zip(rs, expected): |
| self.assertEqual(tuple(r), ex[: len(r)]) |
| |
| # Test boundary values greater than JVM's Long.MaxValue and make sure we don't overflow |
| frame_end = Window.unboundedFollowing + 1 |
| sel = df.select( |
| df.key, F.sum(df.value).over(Window.rowsBetween(Window.currentRow, frame_end)) |
| ) |
| rs = sorted(sel.collect()) |
| expected = [("one", 3), ("two", 2)] |
| for r, ex in zip(rs, expected): |
| self.assertEqual(tuple(r), ex[: len(r)]) |
| |
| def test_collect_functions(self): |
| df = self.spark.createDataFrame([(1, "1"), (2, "2"), (1, "2"), (1, "2")], ["key", "value"]) |
| from pyspark.sql import functions |
| |
| self.assertEqual( |
| sorted(df.select(functions.collect_set(df.key).alias("r")).collect()[0].r), [1, 2] |
| ) |
| self.assertEqual( |
| sorted(df.select(functions.collect_list(df.key).alias("r")).collect()[0].r), |
| [1, 1, 1, 2], |
| ) |
| self.assertEqual( |
| sorted(df.select(functions.collect_set(df.value).alias("r")).collect()[0].r), ["1", "2"] |
| ) |
| self.assertEqual( |
| sorted(df.select(functions.collect_list(df.value).alias("r")).collect()[0].r), |
| ["1", "2", "2", "2"], |
| ) |
| |
| def test_datetime_functions(self): |
| from pyspark.sql import functions |
| from datetime import date |
| |
| df = self.spark.range(1).selectExpr("'2017-01-22' as dateCol") |
| parse_result = df.select(functions.to_date(functions.col("dateCol"))).first() |
| self.assertEqual(date(2017, 1, 22), parse_result["to_date(dateCol)"]) |
| |
| def test_assert_true(self): |
| from pyspark.sql.functions import assert_true |
| |
| df = self.spark.range(3) |
| |
| self.assertEqual( |
| df.select(assert_true(df.id < 3)).toDF("val").collect(), |
| [Row(val=None), Row(val=None), Row(val=None)], |
| ) |
| |
| with self.assertRaises(Py4JJavaError) as cm: |
| df.select(assert_true(df.id < 2, "too big")).toDF("val").collect() |
| self.assertIn("java.lang.RuntimeException", str(cm.exception)) |
| self.assertIn("too big", str(cm.exception)) |
| |
| with self.assertRaises(Py4JJavaError) as cm: |
| df.select(assert_true(df.id < 2, df.id * 1e6)).toDF("val").collect() |
| self.assertIn("java.lang.RuntimeException", str(cm.exception)) |
| self.assertIn("2000000", str(cm.exception)) |
| |
| with self.assertRaises(TypeError) as cm: |
| df.select(assert_true(df.id < 2, 5)) |
| self.assertEqual("errMsg should be a Column or a str, got <class 'int'>", str(cm.exception)) |
| |
| def test_raise_error(self): |
| from pyspark.sql.functions import raise_error |
| |
| df = self.spark.createDataFrame([Row(id="foobar")]) |
| |
| with self.assertRaises(Py4JJavaError) as cm: |
| df.select(raise_error(df.id)).collect() |
| self.assertIn("java.lang.RuntimeException", str(cm.exception)) |
| self.assertIn("foobar", str(cm.exception)) |
| |
| with self.assertRaises(Py4JJavaError) as cm: |
| df.select(raise_error("barfoo")).collect() |
| self.assertIn("java.lang.RuntimeException", str(cm.exception)) |
| self.assertIn("barfoo", str(cm.exception)) |
| |
| with self.assertRaises(TypeError) as cm: |
| df.select(raise_error(None)) |
| self.assertEqual( |
| "errMsg should be a Column or a str, got <class 'NoneType'>", str(cm.exception) |
| ) |
| |
| def test_sum_distinct(self): |
| self.spark.range(10).select( |
| assert_true(sum_distinct(col("id")) == sumDistinct(col("id"))) |
| ).collect() |
| |
| def test_shiftleft(self): |
| self.spark.range(10).select( |
| assert_true(shiftLeft(col("id"), 2) == shiftleft(col("id"), 2)) |
| ).collect() |
| |
| def test_shiftright(self): |
| self.spark.range(10).select( |
| assert_true(shiftRight(col("id"), 2) == shiftright(col("id"), 2)) |
| ).collect() |
| |
| def test_shiftrightunsigned(self): |
| self.spark.range(10).select( |
| assert_true(shiftRightUnsigned(col("id"), 2) == shiftrightunsigned(col("id"), 2)) |
| ).collect() |
| |
| def test_lit_day_time_interval(self): |
| td = datetime.timedelta(days=1, hours=12, milliseconds=123) |
| actual = self.spark.range(1).select(lit(td)).first()[0] |
| self.assertEqual(actual, td) |
| |
| |
| if __name__ == "__main__": |
| import unittest |
| from pyspark.sql.tests.test_functions 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) |
