You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by gu...@apache.org on 2018/11/14 06:51:38 UTC
[4/7] spark git commit: [SPARK-26032][PYTHON] Break large
sql/tests.py files into smaller files
http://git-wip-us.apache.org/repos/asf/spark/blob/a7a331df/python/pyspark/sql/tests/test_dataframe.py
----------------------------------------------------------------------
diff --git a/python/pyspark/sql/tests/test_dataframe.py b/python/pyspark/sql/tests/test_dataframe.py
new file mode 100644
index 0000000..eba00b5
--- /dev/null
+++ b/python/pyspark/sql/tests/test_dataframe.py
@@ -0,0 +1,737 @@
+#
+# 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 os
+import pydoc
+import time
+import unittest
+
+from pyspark.sql import SparkSession, Row
+from pyspark.sql.types import *
+from pyspark.sql.utils import AnalysisException, IllegalArgumentException
+from pyspark.testing.sqlutils import ReusedSQLTestCase, SQLTestUtils, have_pyarrow, have_pandas, \
+ pandas_requirement_message, pyarrow_requirement_message
+from pyspark.tests import QuietTest
+
+
+class DataFrameTests(ReusedSQLTestCase):
+
+ def test_range(self):
+ self.assertEqual(self.spark.range(1, 1).count(), 0)
+ self.assertEqual(self.spark.range(1, 0, -1).count(), 1)
+ self.assertEqual(self.spark.range(0, 1 << 40, 1 << 39).count(), 2)
+ self.assertEqual(self.spark.range(-2).count(), 0)
+ self.assertEqual(self.spark.range(3).count(), 3)
+
+ def test_duplicated_column_names(self):
+ df = self.spark.createDataFrame([(1, 2)], ["c", "c"])
+ row = df.select('*').first()
+ self.assertEqual(1, row[0])
+ self.assertEqual(2, row[1])
+ self.assertEqual("Row(c=1, c=2)", str(row))
+ # Cannot access columns
+ self.assertRaises(AnalysisException, lambda: df.select(df[0]).first())
+ self.assertRaises(AnalysisException, lambda: df.select(df.c).first())
+ self.assertRaises(AnalysisException, lambda: df.select(df["c"]).first())
+
+ def test_freqItems(self):
+ vals = [Row(a=1, b=-2.0) if i % 2 == 0 else Row(a=i, b=i * 1.0) for i in range(100)]
+ df = self.sc.parallelize(vals).toDF()
+ items = df.stat.freqItems(("a", "b"), 0.4).collect()[0]
+ self.assertTrue(1 in items[0])
+ self.assertTrue(-2.0 in items[1])
+
+ def test_help_command(self):
+ # Regression test for SPARK-5464
+ rdd = self.sc.parallelize(['{"foo":"bar"}', '{"foo":"baz"}'])
+ df = self.spark.read.json(rdd)
+ # render_doc() reproduces the help() exception without printing output
+ pydoc.render_doc(df)
+ pydoc.render_doc(df.foo)
+ pydoc.render_doc(df.take(1))
+
+ def test_dropna(self):
+ schema = StructType([
+ StructField("name", StringType(), True),
+ StructField("age", IntegerType(), True),
+ StructField("height", DoubleType(), True)])
+
+ # shouldn't drop a non-null row
+ self.assertEqual(self.spark.createDataFrame(
+ [(u'Alice', 50, 80.1)], schema).dropna().count(),
+ 1)
+
+ # dropping rows with a single null value
+ self.assertEqual(self.spark.createDataFrame(
+ [(u'Alice', None, 80.1)], schema).dropna().count(),
+ 0)
+ self.assertEqual(self.spark.createDataFrame(
+ [(u'Alice', None, 80.1)], schema).dropna(how='any').count(),
+ 0)
+
+ # if how = 'all', only drop rows if all values are null
+ self.assertEqual(self.spark.createDataFrame(
+ [(u'Alice', None, 80.1)], schema).dropna(how='all').count(),
+ 1)
+ self.assertEqual(self.spark.createDataFrame(
+ [(None, None, None)], schema).dropna(how='all').count(),
+ 0)
+
+ # how and subset
+ self.assertEqual(self.spark.createDataFrame(
+ [(u'Alice', 50, None)], schema).dropna(how='any', subset=['name', 'age']).count(),
+ 1)
+ self.assertEqual(self.spark.createDataFrame(
+ [(u'Alice', None, None)], schema).dropna(how='any', subset=['name', 'age']).count(),
+ 0)
+
+ # threshold
+ self.assertEqual(self.spark.createDataFrame(
+ [(u'Alice', None, 80.1)], schema).dropna(thresh=2).count(),
+ 1)
+ self.assertEqual(self.spark.createDataFrame(
+ [(u'Alice', None, None)], schema).dropna(thresh=2).count(),
+ 0)
+
+ # threshold and subset
+ self.assertEqual(self.spark.createDataFrame(
+ [(u'Alice', 50, None)], schema).dropna(thresh=2, subset=['name', 'age']).count(),
+ 1)
+ self.assertEqual(self.spark.createDataFrame(
+ [(u'Alice', None, 180.9)], schema).dropna(thresh=2, subset=['name', 'age']).count(),
+ 0)
+
+ # thresh should take precedence over how
+ self.assertEqual(self.spark.createDataFrame(
+ [(u'Alice', 50, None)], schema).dropna(
+ how='any', thresh=2, subset=['name', 'age']).count(),
+ 1)
+
+ def test_fillna(self):
+ schema = StructType([
+ StructField("name", StringType(), True),
+ StructField("age", IntegerType(), True),
+ StructField("height", DoubleType(), True),
+ StructField("spy", BooleanType(), True)])
+
+ # fillna shouldn't change non-null values
+ row = self.spark.createDataFrame([(u'Alice', 10, 80.1, True)], schema).fillna(50).first()
+ self.assertEqual(row.age, 10)
+
+ # fillna with int
+ row = self.spark.createDataFrame([(u'Alice', None, None, None)], schema).fillna(50).first()
+ self.assertEqual(row.age, 50)
+ self.assertEqual(row.height, 50.0)
+
+ # fillna with double
+ row = self.spark.createDataFrame(
+ [(u'Alice', None, None, None)], schema).fillna(50.1).first()
+ self.assertEqual(row.age, 50)
+ self.assertEqual(row.height, 50.1)
+
+ # fillna with bool
+ row = self.spark.createDataFrame(
+ [(u'Alice', None, None, None)], schema).fillna(True).first()
+ self.assertEqual(row.age, None)
+ self.assertEqual(row.spy, True)
+
+ # fillna with string
+ row = self.spark.createDataFrame([(None, None, None, None)], schema).fillna("hello").first()
+ self.assertEqual(row.name, u"hello")
+ self.assertEqual(row.age, None)
+
+ # fillna with subset specified for numeric cols
+ row = self.spark.createDataFrame(
+ [(None, None, None, None)], schema).fillna(50, subset=['name', 'age']).first()
+ self.assertEqual(row.name, None)
+ self.assertEqual(row.age, 50)
+ self.assertEqual(row.height, None)
+ self.assertEqual(row.spy, None)
+
+ # fillna with subset specified for string cols
+ row = self.spark.createDataFrame(
+ [(None, None, None, None)], schema).fillna("haha", subset=['name', 'age']).first()
+ self.assertEqual(row.name, "haha")
+ self.assertEqual(row.age, None)
+ self.assertEqual(row.height, None)
+ self.assertEqual(row.spy, None)
+
+ # fillna with subset specified for bool cols
+ row = self.spark.createDataFrame(
+ [(None, None, None, None)], schema).fillna(True, subset=['name', 'spy']).first()
+ self.assertEqual(row.name, None)
+ self.assertEqual(row.age, None)
+ self.assertEqual(row.height, None)
+ self.assertEqual(row.spy, True)
+
+ # fillna with dictionary for boolean types
+ row = self.spark.createDataFrame([Row(a=None), Row(a=True)]).fillna({"a": True}).first()
+ self.assertEqual(row.a, True)
+
+ def test_repartitionByRange_dataframe(self):
+ schema = StructType([
+ StructField("name", StringType(), True),
+ StructField("age", IntegerType(), True),
+ StructField("height", DoubleType(), True)])
+
+ df1 = self.spark.createDataFrame(
+ [(u'Bob', 27, 66.0), (u'Alice', 10, 10.0), (u'Bob', 10, 66.0)], schema)
+ df2 = self.spark.createDataFrame(
+ [(u'Alice', 10, 10.0), (u'Bob', 10, 66.0), (u'Bob', 27, 66.0)], schema)
+
+ # test repartitionByRange(numPartitions, *cols)
+ df3 = df1.repartitionByRange(2, "name", "age")
+ self.assertEqual(df3.rdd.getNumPartitions(), 2)
+ self.assertEqual(df3.rdd.first(), df2.rdd.first())
+ self.assertEqual(df3.rdd.take(3), df2.rdd.take(3))
+
+ # test repartitionByRange(numPartitions, *cols)
+ df4 = df1.repartitionByRange(3, "name", "age")
+ self.assertEqual(df4.rdd.getNumPartitions(), 3)
+ self.assertEqual(df4.rdd.first(), df2.rdd.first())
+ self.assertEqual(df4.rdd.take(3), df2.rdd.take(3))
+
+ # test repartitionByRange(*cols)
+ df5 = df1.repartitionByRange("name", "age")
+ self.assertEqual(df5.rdd.first(), df2.rdd.first())
+ self.assertEqual(df5.rdd.take(3), df2.rdd.take(3))
+
+ def test_replace(self):
+ schema = StructType([
+ StructField("name", StringType(), True),
+ StructField("age", IntegerType(), True),
+ StructField("height", DoubleType(), True)])
+
+ # replace with int
+ row = self.spark.createDataFrame([(u'Alice', 10, 10.0)], schema).replace(10, 20).first()
+ self.assertEqual(row.age, 20)
+ self.assertEqual(row.height, 20.0)
+
+ # replace with double
+ row = self.spark.createDataFrame(
+ [(u'Alice', 80, 80.0)], schema).replace(80.0, 82.1).first()
+ self.assertEqual(row.age, 82)
+ self.assertEqual(row.height, 82.1)
+
+ # replace with string
+ row = self.spark.createDataFrame(
+ [(u'Alice', 10, 80.1)], schema).replace(u'Alice', u'Ann').first()
+ self.assertEqual(row.name, u"Ann")
+ self.assertEqual(row.age, 10)
+
+ # replace with subset specified by a string of a column name w/ actual change
+ row = self.spark.createDataFrame(
+ [(u'Alice', 10, 80.1)], schema).replace(10, 20, subset='age').first()
+ self.assertEqual(row.age, 20)
+
+ # replace with subset specified by a string of a column name w/o actual change
+ row = self.spark.createDataFrame(
+ [(u'Alice', 10, 80.1)], schema).replace(10, 20, subset='height').first()
+ self.assertEqual(row.age, 10)
+
+ # replace with subset specified with one column replaced, another column not in subset
+ # stays unchanged.
+ row = self.spark.createDataFrame(
+ [(u'Alice', 10, 10.0)], schema).replace(10, 20, subset=['name', 'age']).first()
+ self.assertEqual(row.name, u'Alice')
+ self.assertEqual(row.age, 20)
+ self.assertEqual(row.height, 10.0)
+
+ # replace with subset specified but no column will be replaced
+ row = self.spark.createDataFrame(
+ [(u'Alice', 10, None)], schema).replace(10, 20, subset=['name', 'height']).first()
+ self.assertEqual(row.name, u'Alice')
+ self.assertEqual(row.age, 10)
+ self.assertEqual(row.height, None)
+
+ # replace with lists
+ row = self.spark.createDataFrame(
+ [(u'Alice', 10, 80.1)], schema).replace([u'Alice'], [u'Ann']).first()
+ self.assertTupleEqual(row, (u'Ann', 10, 80.1))
+
+ # replace with dict
+ row = self.spark.createDataFrame(
+ [(u'Alice', 10, 80.1)], schema).replace({10: 11}).first()
+ self.assertTupleEqual(row, (u'Alice', 11, 80.1))
+
+ # test backward compatibility with dummy value
+ dummy_value = 1
+ row = self.spark.createDataFrame(
+ [(u'Alice', 10, 80.1)], schema).replace({'Alice': 'Bob'}, dummy_value).first()
+ self.assertTupleEqual(row, (u'Bob', 10, 80.1))
+
+ # test dict with mixed numerics
+ row = self.spark.createDataFrame(
+ [(u'Alice', 10, 80.1)], schema).replace({10: -10, 80.1: 90.5}).first()
+ self.assertTupleEqual(row, (u'Alice', -10, 90.5))
+
+ # replace with tuples
+ row = self.spark.createDataFrame(
+ [(u'Alice', 10, 80.1)], schema).replace((u'Alice', ), (u'Bob', )).first()
+ self.assertTupleEqual(row, (u'Bob', 10, 80.1))
+
+ # replace multiple columns
+ row = self.spark.createDataFrame(
+ [(u'Alice', 10, 80.0)], schema).replace((10, 80.0), (20, 90)).first()
+ self.assertTupleEqual(row, (u'Alice', 20, 90.0))
+
+ # test for mixed numerics
+ row = self.spark.createDataFrame(
+ [(u'Alice', 10, 80.0)], schema).replace((10, 80), (20, 90.5)).first()
+ self.assertTupleEqual(row, (u'Alice', 20, 90.5))
+
+ row = self.spark.createDataFrame(
+ [(u'Alice', 10, 80.0)], schema).replace({10: 20, 80: 90.5}).first()
+ self.assertTupleEqual(row, (u'Alice', 20, 90.5))
+
+ # replace with boolean
+ row = (self
+ .spark.createDataFrame([(u'Alice', 10, 80.0)], schema)
+ .selectExpr("name = 'Bob'", 'age <= 15')
+ .replace(False, True).first())
+ self.assertTupleEqual(row, (True, True))
+
+ # replace string with None and then drop None rows
+ row = self.spark.createDataFrame(
+ [(u'Alice', 10, 80.0)], schema).replace(u'Alice', None).dropna()
+ self.assertEqual(row.count(), 0)
+
+ # replace with number and None
+ row = self.spark.createDataFrame(
+ [(u'Alice', 10, 80.0)], schema).replace([10, 80], [20, None]).first()
+ self.assertTupleEqual(row, (u'Alice', 20, None))
+
+ # should fail if subset is not list, tuple or None
+ with self.assertRaises(ValueError):
+ self.spark.createDataFrame(
+ [(u'Alice', 10, 80.1)], schema).replace({10: 11}, subset=1).first()
+
+ # should fail if to_replace and value have different length
+ with self.assertRaises(ValueError):
+ self.spark.createDataFrame(
+ [(u'Alice', 10, 80.1)], schema).replace(["Alice", "Bob"], ["Eve"]).first()
+
+ # should fail if when received unexpected type
+ with self.assertRaises(ValueError):
+ from datetime import datetime
+ self.spark.createDataFrame(
+ [(u'Alice', 10, 80.1)], schema).replace(datetime.now(), datetime.now()).first()
+
+ # should fail if provided mixed type replacements
+ with self.assertRaises(ValueError):
+ self.spark.createDataFrame(
+ [(u'Alice', 10, 80.1)], schema).replace(["Alice", 10], ["Eve", 20]).first()
+
+ with self.assertRaises(ValueError):
+ self.spark.createDataFrame(
+ [(u'Alice', 10, 80.1)], schema).replace({u"Alice": u"Bob", 10: 20}).first()
+
+ with self.assertRaisesRegexp(
+ TypeError,
+ 'value argument is required when to_replace is not a dictionary.'):
+ self.spark.createDataFrame(
+ [(u'Alice', 10, 80.0)], schema).replace(["Alice", "Bob"]).first()
+
+ def test_with_column_with_existing_name(self):
+ keys = self.df.withColumn("key", self.df.key).select("key").collect()
+ self.assertEqual([r.key for r in keys], list(range(100)))
+
+ # regression test for SPARK-10417
+ def test_column_iterator(self):
+
+ def foo():
+ for x in self.df.key:
+ break
+
+ self.assertRaises(TypeError, foo)
+
+ def test_generic_hints(self):
+ from pyspark.sql import DataFrame
+
+ df1 = self.spark.range(10e10).toDF("id")
+ df2 = self.spark.range(10e10).toDF("id")
+
+ self.assertIsInstance(df1.hint("broadcast"), DataFrame)
+ self.assertIsInstance(df1.hint("broadcast", []), DataFrame)
+
+ # Dummy rules
+ self.assertIsInstance(df1.hint("broadcast", "foo", "bar"), DataFrame)
+ self.assertIsInstance(df1.hint("broadcast", ["foo", "bar"]), DataFrame)
+
+ plan = df1.join(df2.hint("broadcast"), "id")._jdf.queryExecution().executedPlan()
+ self.assertEqual(1, plan.toString().count("BroadcastHashJoin"))
+
+ def test_sample(self):
+ self.assertRaisesRegexp(
+ TypeError,
+ "should be a bool, float and number",
+ lambda: self.spark.range(1).sample())
+
+ self.assertRaises(
+ TypeError,
+ lambda: self.spark.range(1).sample("a"))
+
+ self.assertRaises(
+ TypeError,
+ lambda: self.spark.range(1).sample(seed="abc"))
+
+ self.assertRaises(
+ IllegalArgumentException,
+ lambda: self.spark.range(1).sample(-1.0))
+
+ def test_toDF_with_schema_string(self):
+ data = [Row(key=i, value=str(i)) for i in range(100)]
+ rdd = self.sc.parallelize(data, 5)
+
+ df = rdd.toDF("key: int, value: string")
+ self.assertEqual(df.schema.simpleString(), "struct<key:int,value:string>")
+ self.assertEqual(df.collect(), data)
+
+ # different but compatible field types can be used.
+ df = rdd.toDF("key: string, value: string")
+ self.assertEqual(df.schema.simpleString(), "struct<key:string,value:string>")
+ self.assertEqual(df.collect(), [Row(key=str(i), value=str(i)) for i in range(100)])
+
+ # field names can differ.
+ df = rdd.toDF(" a: int, b: string ")
+ self.assertEqual(df.schema.simpleString(), "struct<a:int,b:string>")
+ self.assertEqual(df.collect(), data)
+
+ # number of fields must match.
+ self.assertRaisesRegexp(Exception, "Length of object",
+ lambda: rdd.toDF("key: int").collect())
+
+ # field types mismatch will cause exception at runtime.
+ self.assertRaisesRegexp(Exception, "FloatType can not accept",
+ lambda: rdd.toDF("key: float, value: string").collect())
+
+ # flat schema values will be wrapped into row.
+ df = rdd.map(lambda row: row.key).toDF("int")
+ self.assertEqual(df.schema.simpleString(), "struct<value:int>")
+ self.assertEqual(df.collect(), [Row(key=i) for i in range(100)])
+
+ # users can use DataType directly instead of data type string.
+ df = rdd.map(lambda row: row.key).toDF(IntegerType())
+ self.assertEqual(df.schema.simpleString(), "struct<value:int>")
+ self.assertEqual(df.collect(), [Row(key=i) for i in range(100)])
+
+ def test_join_without_on(self):
+ df1 = self.spark.range(1).toDF("a")
+ df2 = self.spark.range(1).toDF("b")
+
+ with self.sql_conf({"spark.sql.crossJoin.enabled": False}):
+ self.assertRaises(AnalysisException, lambda: df1.join(df2, how="inner").collect())
+
+ with self.sql_conf({"spark.sql.crossJoin.enabled": True}):
+ actual = df1.join(df2, how="inner").collect()
+ expected = [Row(a=0, b=0)]
+ self.assertEqual(actual, expected)
+
+ # Regression test for invalid join methods when on is None, Spark-14761
+ def test_invalid_join_method(self):
+ df1 = self.spark.createDataFrame([("Alice", 5), ("Bob", 8)], ["name", "age"])
+ df2 = self.spark.createDataFrame([("Alice", 80), ("Bob", 90)], ["name", "height"])
+ self.assertRaises(IllegalArgumentException, lambda: df1.join(df2, how="invalid-join-type"))
+
+ # Cartesian products require cross join syntax
+ def test_require_cross(self):
+
+ df1 = self.spark.createDataFrame([(1, "1")], ("key", "value"))
+ df2 = self.spark.createDataFrame([(1, "1")], ("key", "value"))
+
+ # joins without conditions require cross join syntax
+ self.assertRaises(AnalysisException, lambda: df1.join(df2).collect())
+
+ # works with crossJoin
+ self.assertEqual(1, df1.crossJoin(df2).count())
+
+ def test_cache(self):
+ spark = self.spark
+ with self.tempView("tab1", "tab2"):
+ spark.createDataFrame([(2, 2), (3, 3)]).createOrReplaceTempView("tab1")
+ spark.createDataFrame([(2, 2), (3, 3)]).createOrReplaceTempView("tab2")
+ self.assertFalse(spark.catalog.isCached("tab1"))
+ self.assertFalse(spark.catalog.isCached("tab2"))
+ spark.catalog.cacheTable("tab1")
+ self.assertTrue(spark.catalog.isCached("tab1"))
+ self.assertFalse(spark.catalog.isCached("tab2"))
+ spark.catalog.cacheTable("tab2")
+ spark.catalog.uncacheTable("tab1")
+ self.assertFalse(spark.catalog.isCached("tab1"))
+ self.assertTrue(spark.catalog.isCached("tab2"))
+ spark.catalog.clearCache()
+ self.assertFalse(spark.catalog.isCached("tab1"))
+ self.assertFalse(spark.catalog.isCached("tab2"))
+ self.assertRaisesRegexp(
+ AnalysisException,
+ "does_not_exist",
+ lambda: spark.catalog.isCached("does_not_exist"))
+ self.assertRaisesRegexp(
+ AnalysisException,
+ "does_not_exist",
+ lambda: spark.catalog.cacheTable("does_not_exist"))
+ self.assertRaisesRegexp(
+ AnalysisException,
+ "does_not_exist",
+ lambda: spark.catalog.uncacheTable("does_not_exist"))
+
+ def _to_pandas(self):
+ from datetime import datetime, date
+ schema = StructType().add("a", IntegerType()).add("b", StringType())\
+ .add("c", BooleanType()).add("d", FloatType())\
+ .add("dt", DateType()).add("ts", TimestampType())
+ data = [
+ (1, "foo", True, 3.0, date(1969, 1, 1), datetime(1969, 1, 1, 1, 1, 1)),
+ (2, "foo", True, 5.0, None, None),
+ (3, "bar", False, -1.0, date(2012, 3, 3), datetime(2012, 3, 3, 3, 3, 3)),
+ (4, "bar", False, 6.0, date(2100, 4, 4), datetime(2100, 4, 4, 4, 4, 4)),
+ ]
+ df = self.spark.createDataFrame(data, schema)
+ return df.toPandas()
+
+ @unittest.skipIf(not have_pandas, pandas_requirement_message)
+ def test_to_pandas(self):
+ import numpy as np
+ pdf = self._to_pandas()
+ types = pdf.dtypes
+ self.assertEquals(types[0], np.int32)
+ self.assertEquals(types[1], np.object)
+ self.assertEquals(types[2], np.bool)
+ self.assertEquals(types[3], np.float32)
+ self.assertEquals(types[4], np.object) # datetime.date
+ self.assertEquals(types[5], 'datetime64[ns]')
+
+ @unittest.skipIf(have_pandas, "Required Pandas was found.")
+ def test_to_pandas_required_pandas_not_found(self):
+ with QuietTest(self.sc):
+ with self.assertRaisesRegexp(ImportError, 'Pandas >= .* must be installed'):
+ self._to_pandas()
+
+ @unittest.skipIf(not have_pandas, pandas_requirement_message)
+ def test_to_pandas_avoid_astype(self):
+ import numpy as np
+ schema = StructType().add("a", IntegerType()).add("b", StringType())\
+ .add("c", IntegerType())
+ data = [(1, "foo", 16777220), (None, "bar", None)]
+ df = self.spark.createDataFrame(data, schema)
+ types = df.toPandas().dtypes
+ self.assertEquals(types[0], np.float64) # doesn't convert to np.int32 due to NaN value.
+ self.assertEquals(types[1], np.object)
+ self.assertEquals(types[2], np.float64)
+
+ def test_create_dataframe_from_array_of_long(self):
+ import array
+ data = [Row(longarray=array.array('l', [-9223372036854775808, 0, 9223372036854775807]))]
+ df = self.spark.createDataFrame(data)
+ self.assertEqual(df.first(), Row(longarray=[-9223372036854775808, 0, 9223372036854775807]))
+
+ @unittest.skipIf(not have_pandas, pandas_requirement_message)
+ def test_create_dataframe_from_pandas_with_timestamp(self):
+ import pandas as pd
+ from datetime import datetime
+ pdf = pd.DataFrame({"ts": [datetime(2017, 10, 31, 1, 1, 1)],
+ "d": [pd.Timestamp.now().date()]}, columns=["d", "ts"])
+ # test types are inferred correctly without specifying schema
+ df = self.spark.createDataFrame(pdf)
+ self.assertTrue(isinstance(df.schema['ts'].dataType, TimestampType))
+ self.assertTrue(isinstance(df.schema['d'].dataType, DateType))
+ # test with schema will accept pdf as input
+ df = self.spark.createDataFrame(pdf, schema="d date, ts timestamp")
+ self.assertTrue(isinstance(df.schema['ts'].dataType, TimestampType))
+ self.assertTrue(isinstance(df.schema['d'].dataType, DateType))
+
+ @unittest.skipIf(have_pandas, "Required Pandas was found.")
+ def test_create_dataframe_required_pandas_not_found(self):
+ with QuietTest(self.sc):
+ with self.assertRaisesRegexp(
+ ImportError,
+ "(Pandas >= .* must be installed|No module named '?pandas'?)"):
+ import pandas as pd
+ from datetime import datetime
+ pdf = pd.DataFrame({"ts": [datetime(2017, 10, 31, 1, 1, 1)],
+ "d": [pd.Timestamp.now().date()]})
+ self.spark.createDataFrame(pdf)
+
+ # Regression test for SPARK-23360
+ @unittest.skipIf(not have_pandas, pandas_requirement_message)
+ def test_create_dateframe_from_pandas_with_dst(self):
+ import pandas as pd
+ from datetime import datetime
+
+ pdf = pd.DataFrame({'time': [datetime(2015, 10, 31, 22, 30)]})
+
+ df = self.spark.createDataFrame(pdf)
+ self.assertPandasEqual(pdf, df.toPandas())
+
+ orig_env_tz = os.environ.get('TZ', None)
+ try:
+ tz = 'America/Los_Angeles'
+ os.environ['TZ'] = tz
+ time.tzset()
+ with self.sql_conf({'spark.sql.session.timeZone': tz}):
+ df = self.spark.createDataFrame(pdf)
+ self.assertPandasEqual(pdf, df.toPandas())
+ finally:
+ del os.environ['TZ']
+ if orig_env_tz is not None:
+ os.environ['TZ'] = orig_env_tz
+ time.tzset()
+
+ def test_repr_behaviors(self):
+ import re
+ pattern = re.compile(r'^ *\|', re.MULTILINE)
+ df = self.spark.createDataFrame([(1, "1"), (22222, "22222")], ("key", "value"))
+
+ # test when eager evaluation is enabled and _repr_html_ will not be called
+ with self.sql_conf({"spark.sql.repl.eagerEval.enabled": True}):
+ expected1 = """+-----+-----+
+ || key|value|
+ |+-----+-----+
+ || 1| 1|
+ ||22222|22222|
+ |+-----+-----+
+ |"""
+ self.assertEquals(re.sub(pattern, '', expected1), df.__repr__())
+ with self.sql_conf({"spark.sql.repl.eagerEval.truncate": 3}):
+ expected2 = """+---+-----+
+ ||key|value|
+ |+---+-----+
+ || 1| 1|
+ ||222| 222|
+ |+---+-----+
+ |"""
+ self.assertEquals(re.sub(pattern, '', expected2), df.__repr__())
+ with self.sql_conf({"spark.sql.repl.eagerEval.maxNumRows": 1}):
+ expected3 = """+---+-----+
+ ||key|value|
+ |+---+-----+
+ || 1| 1|
+ |+---+-----+
+ |only showing top 1 row
+ |"""
+ self.assertEquals(re.sub(pattern, '', expected3), df.__repr__())
+
+ # test when eager evaluation is enabled and _repr_html_ will be called
+ with self.sql_conf({"spark.sql.repl.eagerEval.enabled": True}):
+ expected1 = """<table border='1'>
+ |<tr><th>key</th><th>value</th></tr>
+ |<tr><td>1</td><td>1</td></tr>
+ |<tr><td>22222</td><td>22222</td></tr>
+ |</table>
+ |"""
+ self.assertEquals(re.sub(pattern, '', expected1), df._repr_html_())
+ with self.sql_conf({"spark.sql.repl.eagerEval.truncate": 3}):
+ expected2 = """<table border='1'>
+ |<tr><th>key</th><th>value</th></tr>
+ |<tr><td>1</td><td>1</td></tr>
+ |<tr><td>222</td><td>222</td></tr>
+ |</table>
+ |"""
+ self.assertEquals(re.sub(pattern, '', expected2), df._repr_html_())
+ with self.sql_conf({"spark.sql.repl.eagerEval.maxNumRows": 1}):
+ expected3 = """<table border='1'>
+ |<tr><th>key</th><th>value</th></tr>
+ |<tr><td>1</td><td>1</td></tr>
+ |</table>
+ |only showing top 1 row
+ |"""
+ self.assertEquals(re.sub(pattern, '', expected3), df._repr_html_())
+
+ # test when eager evaluation is disabled and _repr_html_ will be called
+ with self.sql_conf({"spark.sql.repl.eagerEval.enabled": False}):
+ expected = "DataFrame[key: bigint, value: string]"
+ self.assertEquals(None, df._repr_html_())
+ self.assertEquals(expected, df.__repr__())
+ with self.sql_conf({"spark.sql.repl.eagerEval.truncate": 3}):
+ self.assertEquals(None, df._repr_html_())
+ self.assertEquals(expected, df.__repr__())
+ with self.sql_conf({"spark.sql.repl.eagerEval.maxNumRows": 1}):
+ self.assertEquals(None, df._repr_html_())
+ self.assertEquals(expected, df.__repr__())
+
+
+class QueryExecutionListenerTests(unittest.TestCase, SQLTestUtils):
+ # These tests are separate because it uses 'spark.sql.queryExecutionListeners' which is
+ # static and immutable. This can't be set or unset, for example, via `spark.conf`.
+
+ @classmethod
+ def setUpClass(cls):
+ import glob
+ from pyspark.find_spark_home import _find_spark_home
+
+ SPARK_HOME = _find_spark_home()
+ filename_pattern = (
+ "sql/core/target/scala-*/test-classes/org/apache/spark/sql/"
+ "TestQueryExecutionListener.class")
+ cls.has_listener = bool(glob.glob(os.path.join(SPARK_HOME, filename_pattern)))
+
+ if cls.has_listener:
+ # Note that 'spark.sql.queryExecutionListeners' is a static immutable configuration.
+ cls.spark = SparkSession.builder \
+ .master("local[4]") \
+ .appName(cls.__name__) \
+ .config(
+ "spark.sql.queryExecutionListeners",
+ "org.apache.spark.sql.TestQueryExecutionListener") \
+ .getOrCreate()
+
+ def setUp(self):
+ if not self.has_listener:
+ raise self.skipTest(
+ "'org.apache.spark.sql.TestQueryExecutionListener' is not "
+ "available. Will skip the related tests.")
+
+ @classmethod
+ def tearDownClass(cls):
+ if hasattr(cls, "spark"):
+ cls.spark.stop()
+
+ def tearDown(self):
+ self.spark._jvm.OnSuccessCall.clear()
+
+ def test_query_execution_listener_on_collect(self):
+ self.assertFalse(
+ self.spark._jvm.OnSuccessCall.isCalled(),
+ "The callback from the query execution listener should not be called before 'collect'")
+ self.spark.sql("SELECT * FROM range(1)").collect()
+ self.assertTrue(
+ self.spark._jvm.OnSuccessCall.isCalled(),
+ "The callback from the query execution listener should be called after 'collect'")
+
+ @unittest.skipIf(
+ not have_pandas or not have_pyarrow,
+ pandas_requirement_message or pyarrow_requirement_message)
+ def test_query_execution_listener_on_collect_with_arrow(self):
+ with self.sql_conf({"spark.sql.execution.arrow.enabled": True}):
+ self.assertFalse(
+ self.spark._jvm.OnSuccessCall.isCalled(),
+ "The callback from the query execution listener should not be "
+ "called before 'toPandas'")
+ self.spark.sql("SELECT * FROM range(1)").toPandas()
+ self.assertTrue(
+ self.spark._jvm.OnSuccessCall.isCalled(),
+ "The callback from the query execution listener should be called after 'toPandas'")
+
+
+if __name__ == "__main__":
+ from pyspark.sql.tests.test_dataframe import *
+
+ try:
+ import xmlrunner
+ unittest.main(testRunner=xmlrunner.XMLTestRunner(output='target/test-reports'), verbosity=2)
+ except ImportError:
+ unittest.main(verbosity=2)
http://git-wip-us.apache.org/repos/asf/spark/blob/a7a331df/python/pyspark/sql/tests/test_datasources.py
----------------------------------------------------------------------
diff --git a/python/pyspark/sql/tests/test_datasources.py b/python/pyspark/sql/tests/test_datasources.py
new file mode 100644
index 0000000..b827378
--- /dev/null
+++ b/python/pyspark/sql/tests/test_datasources.py
@@ -0,0 +1,170 @@
+#
+# 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 shutil
+import tempfile
+
+from pyspark.sql import Row
+from pyspark.sql.types import *
+from pyspark.testing.sqlutils import ReusedSQLTestCase
+
+
+class DataSourcesTests(ReusedSQLTestCase):
+
+ def test_linesep_text(self):
+ df = self.spark.read.text("python/test_support/sql/ages_newlines.csv", lineSep=",")
+ expected = [Row(value=u'Joe'), Row(value=u'20'), Row(value=u'"Hi'),
+ Row(value=u'\nI am Jeo"\nTom'), Row(value=u'30'),
+ Row(value=u'"My name is Tom"\nHyukjin'), Row(value=u'25'),
+ Row(value=u'"I am Hyukjin\n\nI love Spark!"\n')]
+ self.assertEqual(df.collect(), expected)
+
+ tpath = tempfile.mkdtemp()
+ shutil.rmtree(tpath)
+ try:
+ df.write.text(tpath, lineSep="!")
+ expected = [Row(value=u'Joe!20!"Hi!'), Row(value=u'I am Jeo"'),
+ Row(value=u'Tom!30!"My name is Tom"'),
+ Row(value=u'Hyukjin!25!"I am Hyukjin'),
+ Row(value=u''), Row(value=u'I love Spark!"'),
+ Row(value=u'!')]
+ readback = self.spark.read.text(tpath)
+ self.assertEqual(readback.collect(), expected)
+ finally:
+ shutil.rmtree(tpath)
+
+ def test_multiline_json(self):
+ people1 = self.spark.read.json("python/test_support/sql/people.json")
+ people_array = self.spark.read.json("python/test_support/sql/people_array.json",
+ multiLine=True)
+ self.assertEqual(people1.collect(), people_array.collect())
+
+ def test_encoding_json(self):
+ people_array = self.spark.read\
+ .json("python/test_support/sql/people_array_utf16le.json",
+ multiLine=True, encoding="UTF-16LE")
+ expected = [Row(age=30, name=u'Andy'), Row(age=19, name=u'Justin')]
+ self.assertEqual(people_array.collect(), expected)
+
+ def test_linesep_json(self):
+ df = self.spark.read.json("python/test_support/sql/people.json", lineSep=",")
+ expected = [Row(_corrupt_record=None, name=u'Michael'),
+ Row(_corrupt_record=u' "age":30}\n{"name":"Justin"', name=None),
+ Row(_corrupt_record=u' "age":19}\n', name=None)]
+ self.assertEqual(df.collect(), expected)
+
+ tpath = tempfile.mkdtemp()
+ shutil.rmtree(tpath)
+ try:
+ df = self.spark.read.json("python/test_support/sql/people.json")
+ df.write.json(tpath, lineSep="!!")
+ readback = self.spark.read.json(tpath, lineSep="!!")
+ self.assertEqual(readback.collect(), df.collect())
+ finally:
+ shutil.rmtree(tpath)
+
+ def test_multiline_csv(self):
+ ages_newlines = self.spark.read.csv(
+ "python/test_support/sql/ages_newlines.csv", multiLine=True)
+ expected = [Row(_c0=u'Joe', _c1=u'20', _c2=u'Hi,\nI am Jeo'),
+ Row(_c0=u'Tom', _c1=u'30', _c2=u'My name is Tom'),
+ Row(_c0=u'Hyukjin', _c1=u'25', _c2=u'I am Hyukjin\n\nI love Spark!')]
+ self.assertEqual(ages_newlines.collect(), expected)
+
+ def test_ignorewhitespace_csv(self):
+ tmpPath = tempfile.mkdtemp()
+ shutil.rmtree(tmpPath)
+ self.spark.createDataFrame([[" a", "b ", " c "]]).write.csv(
+ tmpPath,
+ ignoreLeadingWhiteSpace=False,
+ ignoreTrailingWhiteSpace=False)
+
+ expected = [Row(value=u' a,b , c ')]
+ readback = self.spark.read.text(tmpPath)
+ self.assertEqual(readback.collect(), expected)
+ shutil.rmtree(tmpPath)
+
+ def test_read_multiple_orc_file(self):
+ df = self.spark.read.orc(["python/test_support/sql/orc_partitioned/b=0/c=0",
+ "python/test_support/sql/orc_partitioned/b=1/c=1"])
+ self.assertEqual(2, df.count())
+
+ def test_read_text_file_list(self):
+ df = self.spark.read.text(['python/test_support/sql/text-test.txt',
+ 'python/test_support/sql/text-test.txt'])
+ count = df.count()
+ self.assertEquals(count, 4)
+
+ def test_json_sampling_ratio(self):
+ rdd = self.spark.sparkContext.range(0, 100, 1, 1) \
+ .map(lambda x: '{"a":0.1}' if x == 1 else '{"a":%s}' % str(x))
+ schema = self.spark.read.option('inferSchema', True) \
+ .option('samplingRatio', 0.5) \
+ .json(rdd).schema
+ self.assertEquals(schema, StructType([StructField("a", LongType(), True)]))
+
+ def test_csv_sampling_ratio(self):
+ rdd = self.spark.sparkContext.range(0, 100, 1, 1) \
+ .map(lambda x: '0.1' if x == 1 else str(x))
+ schema = self.spark.read.option('inferSchema', True)\
+ .csv(rdd, samplingRatio=0.5).schema
+ self.assertEquals(schema, StructType([StructField("_c0", IntegerType(), True)]))
+
+ def test_checking_csv_header(self):
+ path = tempfile.mkdtemp()
+ shutil.rmtree(path)
+ try:
+ self.spark.createDataFrame([[1, 1000], [2000, 2]])\
+ .toDF('f1', 'f2').write.option("header", "true").csv(path)
+ schema = StructType([
+ StructField('f2', IntegerType(), nullable=True),
+ StructField('f1', IntegerType(), nullable=True)])
+ df = self.spark.read.option('header', 'true').schema(schema)\
+ .csv(path, enforceSchema=False)
+ self.assertRaisesRegexp(
+ Exception,
+ "CSV header does not conform to the schema",
+ lambda: df.collect())
+ finally:
+ shutil.rmtree(path)
+
+ def test_ignore_column_of_all_nulls(self):
+ path = tempfile.mkdtemp()
+ shutil.rmtree(path)
+ try:
+ df = self.spark.createDataFrame([["""{"a":null, "b":1, "c":3.0}"""],
+ ["""{"a":null, "b":null, "c":"string"}"""],
+ ["""{"a":null, "b":null, "c":null}"""]])
+ df.write.text(path)
+ schema = StructType([
+ StructField('b', LongType(), nullable=True),
+ StructField('c', StringType(), nullable=True)])
+ readback = self.spark.read.json(path, dropFieldIfAllNull=True)
+ self.assertEquals(readback.schema, schema)
+ finally:
+ shutil.rmtree(path)
+
+
+if __name__ == "__main__":
+ import unittest
+ from pyspark.sql.tests.test_datasources import *
+
+ try:
+ import xmlrunner
+ unittest.main(testRunner=xmlrunner.XMLTestRunner(output='target/test-reports'), verbosity=2)
+ except ImportError:
+ unittest.main(verbosity=2)
http://git-wip-us.apache.org/repos/asf/spark/blob/a7a331df/python/pyspark/sql/tests/test_functions.py
----------------------------------------------------------------------
diff --git a/python/pyspark/sql/tests/test_functions.py b/python/pyspark/sql/tests/test_functions.py
new file mode 100644
index 0000000..f0b59e8
--- /dev/null
+++ b/python/pyspark/sql/tests/test_functions.py
@@ -0,0 +1,278 @@
+#
+# 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
+import sys
+
+from pyspark.sql import Row
+from pyspark.testing.sqlutils import ReusedSQLTestCase
+
+
+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(u"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(10)]).toDF()
+ sampled = df.stat.sampleBy(u"b", fractions={0: 0.5, 1: 0.5}, seed=0)
+ self.assertTrue(sampled.count() == 3)
+
+ def test_cov(self):
+ df = self.sc.parallelize([Row(a=i, b=2 * i) for i in range(10)]).toDF()
+ cov = df.stat.cov(u"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(u"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
+ import math
+
+ def get_values(l):
+ return [j[0] for j in l]
+
+ def assert_close(a, b):
+ c = get_values(b)
+ diff = [abs(v - c[k]) < 1e-6 for k, v in enumerate(a)]
+ return sum(diff) == len(a)
+ assert_close([math.cos(i) for i in range(10)],
+ df.select(functions.cos(df.a)).collect())
+ assert_close([math.cos(i) for i in range(10)],
+ df.select(functions.cos("a")).collect())
+ assert_close([math.sin(i) for i in range(10)],
+ df.select(functions.sin(df.a)).collect())
+ assert_close([math.sin(i) for i in range(10)],
+ df.select(functions.sin(df['a'])).collect())
+ assert_close([math.pow(i, 2 * i) for i in range(10)],
+ df.select(functions.pow(df.a, df.b)).collect())
+ assert_close([math.pow(i, 2) for i in range(10)],
+ df.select(functions.pow(df.a, 2)).collect())
+ assert_close([math.pow(i, 2) for i in range(10)],
+ df.select(functions.pow(df.a, 2.0)).collect())
+ assert_close([math.hypot(i, 2 * i) for i in range(10)],
+ df.select(functions.hypot(df.a, df.b)).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.functions import col, lit
+ df = self.spark.createDataFrame([['nick']], schema=['name'])
+ self.assertRaisesRegexp(
+ TypeError,
+ "must be the same type",
+ lambda: df.select(col('name').substr(0, lit(1))))
+ if sys.version_info.major == 2:
+ self.assertRaises(
+ TypeError,
+ lambda: df.select(col('name').substr(long(0), long(1))))
+
+ 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)
+
+ 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", u"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", u"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((u"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(ValueError, lambda: df.stat.approxQuantile(123, [0.1, 0.9], 0.1))
+ self.assertRaises(ValueError, lambda: df.stat.approxQuantile(("a", 123), [0.1, 0.9], 0.1))
+ self.assertRaises(ValueError, lambda: df.stat.approxQuantile(["a", 123], [0.1, 0.9], 0.1))
+
+ 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.assertEquals(
+ df.select(df.name).orderBy(functions.asc_nulls_first('name')).collect(),
+ [Row(name=None), Row(name=u'Alice'), Row(name=u'Tom')])
+ self.assertEquals(
+ df.select(df.name).orderBy(functions.asc_nulls_last('name')).collect(),
+ [Row(name=u'Alice'), Row(name=u'Tom'), Row(name=None)])
+ self.assertEquals(
+ df.select(df.name).orderBy(functions.desc_nulls_first('name')).collect(),
+ [Row(name=None), Row(name=u'Tom'), Row(name=u'Alice')])
+ self.assertEquals(
+ df.select(df.name).orderBy(functions.desc_nulls_last('name')).collect(),
+ [Row(name=u'Tom'), Row(name=u'Alice'), Row(name=None)])
+
+
+if __name__ == "__main__":
+ import unittest
+ from pyspark.sql.tests.test_functions import *
+
+ try:
+ import xmlrunner
+ unittest.main(testRunner=xmlrunner.XMLTestRunner(output='target/test-reports'), verbosity=2)
+ except ImportError:
+ unittest.main(verbosity=2)
http://git-wip-us.apache.org/repos/asf/spark/blob/a7a331df/python/pyspark/sql/tests/test_group.py
----------------------------------------------------------------------
diff --git a/python/pyspark/sql/tests/test_group.py b/python/pyspark/sql/tests/test_group.py
new file mode 100644
index 0000000..076899f
--- /dev/null
+++ b/python/pyspark/sql/tests/test_group.py
@@ -0,0 +1,45 @@
+#
+# 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 import Row
+from pyspark.testing.sqlutils import ReusedSQLTestCase
+
+
+class GroupTests(ReusedSQLTestCase):
+
+ def test_aggregator(self):
+ df = self.df
+ g = df.groupBy()
+ self.assertEqual([99, 100], sorted(g.agg({'key': 'max', 'value': 'count'}).collect()[0]))
+ self.assertEqual([Row(**{"AVG(key#0)": 49.5})], g.mean().collect())
+
+ from pyspark.sql import functions
+ self.assertEqual((0, u'99'),
+ tuple(g.agg(functions.first(df.key), functions.last(df.value)).first()))
+ self.assertTrue(95 < g.agg(functions.approx_count_distinct(df.key)).first()[0])
+ self.assertEqual(100, g.agg(functions.countDistinct(df.value)).first()[0])
+
+
+if __name__ == "__main__":
+ import unittest
+ from pyspark.sql.tests.test_group import *
+
+ try:
+ import xmlrunner
+ unittest.main(testRunner=xmlrunner.XMLTestRunner(output='target/test-reports'), verbosity=2)
+ except ImportError:
+ unittest.main(verbosity=2)
http://git-wip-us.apache.org/repos/asf/spark/blob/a7a331df/python/pyspark/sql/tests/test_pandas_udf.py
----------------------------------------------------------------------
diff --git a/python/pyspark/sql/tests/test_pandas_udf.py b/python/pyspark/sql/tests/test_pandas_udf.py
new file mode 100644
index 0000000..54a34a7
--- /dev/null
+++ b/python/pyspark/sql/tests/test_pandas_udf.py
@@ -0,0 +1,216 @@
+#
+# 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.sql.types import *
+from pyspark.sql.utils import ParseException
+from pyspark.testing.sqlutils import ReusedSQLTestCase, have_pandas, have_pyarrow, \
+ pandas_requirement_message, pyarrow_requirement_message
+from pyspark.tests import QuietTest
+
+
+@unittest.skipIf(
+ not have_pandas or not have_pyarrow,
+ pandas_requirement_message or pyarrow_requirement_message)
+class PandasUDFTests(ReusedSQLTestCase):
+
+ def test_pandas_udf_basic(self):
+ from pyspark.rdd import PythonEvalType
+ from pyspark.sql.functions import pandas_udf, PandasUDFType
+
+ udf = pandas_udf(lambda x: x, DoubleType())
+ self.assertEqual(udf.returnType, DoubleType())
+ self.assertEqual(udf.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF)
+
+ udf = pandas_udf(lambda x: x, DoubleType(), PandasUDFType.SCALAR)
+ self.assertEqual(udf.returnType, DoubleType())
+ self.assertEqual(udf.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF)
+
+ udf = pandas_udf(lambda x: x, 'double', PandasUDFType.SCALAR)
+ self.assertEqual(udf.returnType, DoubleType())
+ self.assertEqual(udf.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF)
+
+ udf = pandas_udf(lambda x: x, StructType([StructField("v", DoubleType())]),
+ PandasUDFType.GROUPED_MAP)
+ self.assertEqual(udf.returnType, StructType([StructField("v", DoubleType())]))
+ self.assertEqual(udf.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF)
+
+ udf = pandas_udf(lambda x: x, 'v double', PandasUDFType.GROUPED_MAP)
+ self.assertEqual(udf.returnType, StructType([StructField("v", DoubleType())]))
+ self.assertEqual(udf.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF)
+
+ udf = pandas_udf(lambda x: x, 'v double',
+ functionType=PandasUDFType.GROUPED_MAP)
+ self.assertEqual(udf.returnType, StructType([StructField("v", DoubleType())]))
+ self.assertEqual(udf.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF)
+
+ udf = pandas_udf(lambda x: x, returnType='v double',
+ functionType=PandasUDFType.GROUPED_MAP)
+ self.assertEqual(udf.returnType, StructType([StructField("v", DoubleType())]))
+ self.assertEqual(udf.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF)
+
+ def test_pandas_udf_decorator(self):
+ from pyspark.rdd import PythonEvalType
+ from pyspark.sql.functions import pandas_udf, PandasUDFType
+ from pyspark.sql.types import StructType, StructField, DoubleType
+
+ @pandas_udf(DoubleType())
+ def foo(x):
+ return x
+ self.assertEqual(foo.returnType, DoubleType())
+ self.assertEqual(foo.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF)
+
+ @pandas_udf(returnType=DoubleType())
+ def foo(x):
+ return x
+ self.assertEqual(foo.returnType, DoubleType())
+ self.assertEqual(foo.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF)
+
+ schema = StructType([StructField("v", DoubleType())])
+
+ @pandas_udf(schema, PandasUDFType.GROUPED_MAP)
+ def foo(x):
+ return x
+ self.assertEqual(foo.returnType, schema)
+ self.assertEqual(foo.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF)
+
+ @pandas_udf('v double', PandasUDFType.GROUPED_MAP)
+ def foo(x):
+ return x
+ self.assertEqual(foo.returnType, schema)
+ self.assertEqual(foo.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF)
+
+ @pandas_udf(schema, functionType=PandasUDFType.GROUPED_MAP)
+ def foo(x):
+ return x
+ self.assertEqual(foo.returnType, schema)
+ self.assertEqual(foo.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF)
+
+ @pandas_udf(returnType='double', functionType=PandasUDFType.SCALAR)
+ def foo(x):
+ return x
+ self.assertEqual(foo.returnType, DoubleType())
+ self.assertEqual(foo.evalType, PythonEvalType.SQL_SCALAR_PANDAS_UDF)
+
+ @pandas_udf(returnType=schema, functionType=PandasUDFType.GROUPED_MAP)
+ def foo(x):
+ return x
+ self.assertEqual(foo.returnType, schema)
+ self.assertEqual(foo.evalType, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF)
+
+ def test_udf_wrong_arg(self):
+ from pyspark.sql.functions import pandas_udf, PandasUDFType
+
+ with QuietTest(self.sc):
+ with self.assertRaises(ParseException):
+ @pandas_udf('blah')
+ def foo(x):
+ return x
+ with self.assertRaisesRegexp(ValueError, 'Invalid returnType.*None'):
+ @pandas_udf(functionType=PandasUDFType.SCALAR)
+ def foo(x):
+ return x
+ with self.assertRaisesRegexp(ValueError, 'Invalid functionType'):
+ @pandas_udf('double', 100)
+ def foo(x):
+ return x
+
+ with self.assertRaisesRegexp(ValueError, '0-arg pandas_udfs.*not.*supported'):
+ pandas_udf(lambda: 1, LongType(), PandasUDFType.SCALAR)
+ with self.assertRaisesRegexp(ValueError, '0-arg pandas_udfs.*not.*supported'):
+ @pandas_udf(LongType(), PandasUDFType.SCALAR)
+ def zero_with_type():
+ return 1
+
+ with self.assertRaisesRegexp(TypeError, 'Invalid returnType'):
+ @pandas_udf(returnType=PandasUDFType.GROUPED_MAP)
+ def foo(df):
+ return df
+ with self.assertRaisesRegexp(TypeError, 'Invalid returnType'):
+ @pandas_udf(returnType='double', functionType=PandasUDFType.GROUPED_MAP)
+ def foo(df):
+ return df
+ with self.assertRaisesRegexp(ValueError, 'Invalid function'):
+ @pandas_udf(returnType='k int, v double', functionType=PandasUDFType.GROUPED_MAP)
+ def foo(k, v, w):
+ return k
+
+ def test_stopiteration_in_udf(self):
+ from pyspark.sql.functions import udf, pandas_udf, PandasUDFType
+ from py4j.protocol import Py4JJavaError
+
+ def foo(x):
+ raise StopIteration()
+
+ def foofoo(x, y):
+ raise StopIteration()
+
+ exc_message = "Caught StopIteration thrown from user's code; failing the task"
+ df = self.spark.range(0, 100)
+
+ # plain udf (test for SPARK-23754)
+ self.assertRaisesRegexp(
+ Py4JJavaError,
+ exc_message,
+ df.withColumn('v', udf(foo)('id')).collect
+ )
+
+ # pandas scalar udf
+ self.assertRaisesRegexp(
+ Py4JJavaError,
+ exc_message,
+ df.withColumn(
+ 'v', pandas_udf(foo, 'double', PandasUDFType.SCALAR)('id')
+ ).collect
+ )
+
+ # pandas grouped map
+ self.assertRaisesRegexp(
+ Py4JJavaError,
+ exc_message,
+ df.groupBy('id').apply(
+ pandas_udf(foo, df.schema, PandasUDFType.GROUPED_MAP)
+ ).collect
+ )
+
+ self.assertRaisesRegexp(
+ Py4JJavaError,
+ exc_message,
+ df.groupBy('id').apply(
+ pandas_udf(foofoo, df.schema, PandasUDFType.GROUPED_MAP)
+ ).collect
+ )
+
+ # pandas grouped agg
+ self.assertRaisesRegexp(
+ Py4JJavaError,
+ exc_message,
+ df.groupBy('id').agg(
+ pandas_udf(foo, 'double', PandasUDFType.GROUPED_AGG)('id')
+ ).collect
+ )
+
+
+if __name__ == "__main__":
+ from pyspark.sql.tests.test_pandas_udf import *
+
+ try:
+ import xmlrunner
+ unittest.main(testRunner=xmlrunner.XMLTestRunner(output='target/test-reports'), verbosity=2)
+ except ImportError:
+ unittest.main(verbosity=2)
http://git-wip-us.apache.org/repos/asf/spark/blob/a7a331df/python/pyspark/sql/tests/test_pandas_udf_grouped_agg.py
----------------------------------------------------------------------
diff --git a/python/pyspark/sql/tests/test_pandas_udf_grouped_agg.py b/python/pyspark/sql/tests/test_pandas_udf_grouped_agg.py
new file mode 100644
index 0000000..bca47cc
--- /dev/null
+++ b/python/pyspark/sql/tests/test_pandas_udf_grouped_agg.py
@@ -0,0 +1,503 @@
+#
+# 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.sql.types import *
+from pyspark.sql.utils import AnalysisException
+from pyspark.testing.sqlutils import ReusedSQLTestCase, have_pandas, have_pyarrow, \
+ pandas_requirement_message, pyarrow_requirement_message
+from pyspark.tests import QuietTest
+
+
+@unittest.skipIf(
+ not have_pandas or not have_pyarrow,
+ pandas_requirement_message or pyarrow_requirement_message)
+class GroupedAggPandasUDFTests(ReusedSQLTestCase):
+
+ @property
+ def data(self):
+ from pyspark.sql.functions import array, explode, col, lit
+ 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):
+ from pyspark.sql.functions import udf
+
+ @udf('double')
+ def plus_one(v):
+ assert isinstance(v, (int, float))
+ return v + 1
+ return plus_one
+
+ @property
+ def pandas_scalar_plus_two(self):
+ import pandas as pd
+ from pyspark.sql.functions import pandas_udf, PandasUDFType
+
+ @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):
+ from pyspark.sql.functions import pandas_udf, PandasUDFType
+
+ @pandas_udf('double', PandasUDFType.GROUPED_AGG)
+ def avg(v):
+ return v.mean()
+ return avg
+
+ @property
+ def pandas_agg_sum_udf(self):
+ from pyspark.sql.functions import pandas_udf, PandasUDFType
+
+ @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
+ from pyspark.sql.functions import pandas_udf, PandasUDFType
+
+ @pandas_udf('double', PandasUDFType.GROUPED_AGG)
+ def weighted_mean(v, w):
+ return np.average(v, weights=w)
+ return weighted_mean
+
+ def test_manual(self):
+ from pyspark.sql.functions import pandas_udf, array
+
+ 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))'])
+
+ self.assertPandasEqual(expected1.toPandas(), result1.toPandas())
+
+ def test_basic(self):
+ from pyspark.sql.functions import col, lit, mean
+
+ 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')
+ self.assertPandasEqual(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)
+ self.assertPandasEqual(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')
+ self.assertPandasEqual(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')
+ self.assertPandasEqual(expected4.toPandas(), result4.toPandas())
+
+ def test_unsupported_types(self):
+ from pyspark.sql.types import DoubleType, MapType
+ from pyspark.sql.functions import pandas_udf, PandasUDFType
+
+ with QuietTest(self.sc):
+ with self.assertRaisesRegexp(NotImplementedError, 'not supported'):
+ pandas_udf(
+ lambda x: x,
+ ArrayType(ArrayType(TimestampType())),
+ PandasUDFType.GROUPED_AGG)
+
+ with QuietTest(self.sc):
+ with self.assertRaisesRegexp(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.assertRaisesRegexp(NotImplementedError, 'not supported'):
+ @pandas_udf(MapType(DoubleType(), DoubleType()), PandasUDFType.GROUPED_AGG)
+ def mean_and_std_udf(v):
+ return {v.mean(): v.std()}
+
+ def test_alias(self):
+ from pyspark.sql.functions import mean
+
+ 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'))
+
+ self.assertPandasEqual(expected1.toPandas(), result1.toPandas())
+
+ def test_mixed_sql(self):
+ """
+ Test mixing group aggregate pandas UDF with sql expression.
+ """
+ from pyspark.sql.functions import sum
+
+ 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'))
+
+ self.assertPandasEqual(expected1.toPandas(), result1.toPandas())
+ self.assertPandasEqual(expected2.toPandas(), result2.toPandas())
+ self.assertPandasEqual(expected3.toPandas(), result3.toPandas())
+
+ def test_mixed_udfs(self):
+ """
+ Test mixing group aggregate pandas UDF with python UDF and scalar pandas UDF.
+ """
+ from pyspark.sql.functions import sum
+
+ 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)'))
+
+ self.assertPandasEqual(expected1.toPandas(), result1.toPandas())
+ self.assertPandasEqual(expected2.toPandas(), result2.toPandas())
+ self.assertPandasEqual(expected3.toPandas(), result3.toPandas())
+ self.assertPandasEqual(expected4.toPandas(), result4.toPandas())
+ self.assertPandasEqual(expected5.toPandas(), result5.toPandas())
+ self.assertPandasEqual(expected6.toPandas(), result6.toPandas())
+
+ def test_multiple_udfs(self):
+ """
+ Test multiple group aggregate pandas UDFs in one agg function.
+ """
+ from pyspark.sql.functions import sum, mean
+
+ 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())
+
+ self.assertPandasEqual(expected1, result1)
+
+ def test_complex_groupby(self):
+ from pyspark.sql.functions import sum
+
+ 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))
+ expected4 = df.groupby(plus_one(df.id)).agg(sum(df.v))
+
+ # groupby one scalar pandas UDF
+ result5 = df.groupby(plus_two(df.id)).agg(sum_udf(df.v))
+ expected5 = df.groupby(plus_two(df.id)).agg(sum(df.v))
+
+ # groupby one expression and one python UDF
+ result6 = df.groupby(df.v % 2, plus_one(df.id)).agg(sum_udf(df.v))
+ expected6 = df.groupby(df.v % 2, plus_one(df.id)).agg(sum(df.v))
+
+ # 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)')
+ expected7 = df.groupby(df.v % 2, plus_two(df.id)).agg(sum(df.v)).sort('sum(v)')
+
+ self.assertPandasEqual(expected1.toPandas(), result1.toPandas())
+ self.assertPandasEqual(expected2.toPandas(), result2.toPandas())
+ self.assertPandasEqual(expected3.toPandas(), result3.toPandas())
+ self.assertPandasEqual(expected4.toPandas(), result4.toPandas())
+ self.assertPandasEqual(expected5.toPandas(), result5.toPandas())
+ self.assertPandasEqual(expected6.toPandas(), result6.toPandas())
+ self.assertPandasEqual(expected7.toPandas(), result7.toPandas())
+
+ def test_complex_expressions(self):
+ from pyspark.sql.functions import col, sum
+
+ 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')
+ .toPandas())
+
+ 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')
+ .toPandas())
+
+ # 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')
+ .toPandas())
+
+ 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')
+ .toPandas())
+
+ # 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())
+
+ self.assertPandasEqual(expected1, result1)
+ self.assertPandasEqual(expected2, result2)
+ self.assertPandasEqual(expected3, result3)
+
+ def test_retain_group_columns(self):
+ from pyspark.sql.functions import sum
+ 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))
+ self.assertPandasEqual(expected1.toPandas(), result1.toPandas())
+
+ def test_array_type(self):
+ from pyspark.sql.functions import pandas_udf, PandasUDFType
+
+ 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.assertEquals(result1.first()['v2'], [1.0, 2.0])
+
+ def test_invalid_args(self):
+ from pyspark.sql.functions import mean
+
+ df = self.data
+ plus_one = self.python_plus_one
+ mean_udf = self.pandas_agg_mean_udf
+
+ with QuietTest(self.sc):
+ with self.assertRaisesRegexp(
+ AnalysisException,
+ 'nor.*aggregate function'):
+ df.groupby(df.id).agg(plus_one(df.v)).collect()
+
+ with QuietTest(self.sc):
+ with self.assertRaisesRegexp(
+ AnalysisException,
+ 'aggregate function.*argument.*aggregate function'):
+ df.groupby(df.id).agg(mean_udf(mean_udf(df.v))).collect()
+
+ with QuietTest(self.sc):
+ with self.assertRaisesRegexp(
+ 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):
+ from pyspark.sql.functions import pandas_udf
+ from pyspark.rdd import PythonEvalType
+
+ 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)
+
+
+if __name__ == "__main__":
+ from pyspark.sql.tests.test_pandas_udf_grouped_agg import *
+
+ try:
+ import xmlrunner
+ unittest.main(testRunner=xmlrunner.XMLTestRunner(output='target/test-reports'), verbosity=2)
+ except ImportError:
+ unittest.main(verbosity=2)
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org