# -*- encoding: utf-8 -*-

#

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

#

"""

Unit tests for pyspark.sql; additional tests are implemented as doctests in

individual modules.

"""

import os

import sys

import subprocess

import pydoc

import shutil

import tempfile

import pickle

import functools

import time

import datetime

import py4j

try:

import xmlrunner

except ImportError:

xmlrunner = None

if sys.version_info[:2] <= (2, 6):

try:

import unittest2 as unittest

except ImportError:

sys.stderr.write('Please install unittest2 to test with Python 2.6 or earlier')

sys.exit(1)

else:

import unittest

from pyspark.sql import SparkSession, HiveContext, Column, Row

from pyspark.sql.types import *

from pyspark.sql.types import UserDefinedType, _infer_type

from pyspark.tests import ReusedPySparkTestCase, SparkSubmitTests

from pyspark.sql.functions import UserDefinedFunction, sha2

from pyspark.sql.window import Window

from pyspark.sql.utils import AnalysisException, ParseException, IllegalArgumentException

class UTCOffsetTimezone(datetime.tzinfo):

"""

Specifies timezone in UTC offset

"""

def __init__(self, offset=0):

self.ZERO = datetime.timedelta(hours=offset)

def utcoffset(self, dt):

return self.ZERO

def dst(self, dt):

return self.ZERO

class ExamplePointUDT(UserDefinedType):

"""

User-defined type (UDT) for ExamplePoint.

"""

@classmethod

def sqlType(self):

return ArrayType(DoubleType(), False)

@classmethod

def module(cls):

return 'pyspark.sql.tests'

@classmethod

def scalaUDT(cls):

return 'org.apache.spark.sql.test.ExamplePointUDT'

def serialize(self, obj):

return [obj.x, obj.y]

def deserialize(self, datum):

return ExamplePoint(datum[0], datum[1])

class ExamplePoint:

"""

An example class to demonstrate UDT in Scala, Java, and Python.

"""

__UDT__ = ExamplePointUDT()

def __init__(self, x, y):

self.x = x

self.y = y

def __repr__(self):

return "ExamplePoint(%s,%s)" % (self.x, self.y)

def __str__(self):

return "(%s,%s)" % (self.x, self.y)

def __eq__(self, other):

return isinstance(other, self.__class__) and \

other.x == self.x and other.y == self.y

class PythonOnlyUDT(UserDefinedType):

"""

User-defined type (UDT) for ExamplePoint.

"""

@classmethod

def sqlType(self):

return ArrayType(DoubleType(), False)

@classmethod

def module(cls):

return '__main__'

def serialize(self, obj):

return [obj.x, obj.y]

def deserialize(self, datum):

return PythonOnlyPoint(datum[0], datum[1])

@staticmethod

def foo():

pass

@property

def props(self):

return {}

class PythonOnlyPoint(ExamplePoint):

"""

An example class to demonstrate UDT in only Python

"""

__UDT__ = PythonOnlyUDT()

class MyObject(object):

def __init__(self, key, value):

self.key = key

self.value = value

class DataTypeTests(unittest.TestCase):

# regression test for SPARK-6055

def test_data_type_eq(self):

lt = LongType()

lt2 = pickle.loads(pickle.dumps(LongType()))

self.assertEqual(lt, lt2)

# regression test for SPARK-7978

def test_decimal_type(self):

t1 = DecimalType()

t2 = DecimalType(10, 2)

self.assertTrue(t2 is not t1)

self.assertNotEqual(t1, t2)

t3 = DecimalType(8)

self.assertNotEqual(t2, t3)

# regression test for SPARK-10392

def test_datetype_equal_zero(self):

dt = DateType()

self.assertEqual(dt.fromInternal(0), datetime.date(1970, 1, 1))

# regression test for SPARK-17035

def test_timestamp_microsecond(self):

tst = TimestampType()

self.assertEqual(tst.toInternal(datetime.datetime.max) % 1000000, 999999)

def test_empty_row(self):

row = Row()

self.assertEqual(len(row), 0)

class SQLTests(ReusedPySparkTestCase):

@classmethod

def setUpClass(cls):

ReusedPySparkTestCase.setUpClass()

cls.tempdir = tempfile.NamedTemporaryFile(delete=False)

os.unlink(cls.tempdir.name)

cls.spark = SparkSession(cls.sc)

cls.testData = [Row(key=i, value=str(i)) for i in range(100)]

cls.df = cls.spark.createDataFrame(cls.testData)

@classmethod

def tearDownClass(cls):

ReusedPySparkTestCase.tearDownClass()

cls.spark.stop()

shutil.rmtree(cls.tempdir.name, ignore_errors=True)

def test_row_should_be_read_only(self):

row = Row(a=1, b=2)

self.assertEqual(1, row.a)

def foo():

row.a = 3

self.assertRaises(Exception, foo)

row2 = self.spark.range(10).first()

self.assertEqual(0, row2.id)

def foo2():

row2.id = 2

self.assertRaises(Exception, foo2)

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_column_name_encoding(self):

"""Ensure that created columns has `str` type consistently."""

columns = self.spark.createDataFrame([('Alice', 1)], ['name', u'age']).columns

self.assertEqual(columns, ['name', 'age'])

self.assertTrue(isinstance(columns[0], str))

self.assertTrue(isinstance(columns[1], str))

def test_explode(self):

from pyspark.sql.functions import explode

d = [Row(a=1, intlist=[1, 2, 3], mapfield={"a": "b"})]

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")

def test_and_in_expression(self):

self.assertEqual(4, self.df.filter((self.df.key <= 10) & (self.df.value <= "2")).count())

self.assertRaises(ValueError, lambda: (self.df.key <= 10) and (self.df.value <= "2"))

self.assertEqual(14, self.df.filter((self.df.key <= 3) | (self.df.value < "2")).count())

self.assertRaises(ValueError, lambda: self.df.key <= 3 or self.df.value < "2")

self.assertEqual(99, self.df.filter(~(self.df.key == 1)).count())

self.assertRaises(ValueError, lambda: not self.df.key == 1)

def test_udf_with_callable(self):

d = [Row(number=i, squared=i**2) for i in range(10)]

rdd = self.sc.parallelize(d)

data = self.spark.createDataFrame(rdd)

class PlusFour:

def __call__(self, col):

if col is not None:

return col + 4

call = PlusFour()

pudf = UserDefinedFunction(call, LongType())

res = data.select(pudf(data['number']).alias('plus_four'))

self.assertEqual(res.agg({'plus_four': 'sum'}).collect()[0][0], 85)

def test_udf_with_partial_function(self):

d = [Row(number=i, squared=i**2) for i in range(10)]

rdd = self.sc.parallelize(d)

data = self.spark.createDataFrame(rdd)

def some_func(col, param):

if col is not None:

return col + param

pfunc = functools.partial(some_func, param=4)

pudf = UserDefinedFunction(pfunc, LongType())

res = data.select(pudf(data['number']).alias('plus_four'))

self.assertEqual(res.agg({'plus_four': 'sum'}).collect()[0][0], 85)

def test_udf(self):

self.spark.catalog.registerFunction("twoArgs", lambda x, y: len(x) + y, IntegerType())

[row] = self.spark.sql("SELECT twoArgs('test', 1)").collect()

self.assertEqual(row[0], 5)

def test_udf2(self):

self.spark.catalog.registerFunction("strlen", lambda string: len(string), IntegerType())

self.spark.createDataFrame(self.sc.parallelize([Row(a="test")]))\

.createOrReplaceTempView("test")

[res] = self.spark.sql("SELECT strlen(a) FROM test WHERE strlen(a) > 1").collect()

self.assertEqual(4, res[0])

def test_chained_udf(self):

self.spark.catalog.registerFunction("double", lambda x: x + x, IntegerType())

[row] = self.spark.sql("SELECT double(1)").collect()

self.assertEqual(row[0], 2)

[row] = self.spark.sql("SELECT double(double(1))").collect()

self.assertEqual(row[0], 4)

[row] = self.spark.sql("SELECT double(double(1) + 1)").collect()

self.assertEqual(row[0], 6)

def test_multiple_udfs(self):

self.spark.catalog.registerFunction("double", lambda x: x * 2, IntegerType())

[row] = self.spark.sql("SELECT double(1), double(2)").collect()

self.assertEqual(tuple(row), (2, 4))

[row] = self.spark.sql("SELECT double(double(1)), double(double(2) + 2)").collect()

self.assertEqual(tuple(row), (4, 12))

self.spark.catalog.registerFunction("add", lambda x, y: x + y, IntegerType())

[row] = self.spark.sql("SELECT double(add(1, 2)), add(double(2), 1)").collect()

self.assertEqual(tuple(row), (6, 5))

def test_udf_in_filter_on_top_of_outer_join(self):

from pyspark.sql.functions import udf

left = self.spark.createDataFrame([Row(a=1)])

right = self.spark.createDataFrame([Row(a=1)])

df = left.join(right, on='a', how='left_outer')

df = df.withColumn('b', udf(lambda x: 'x')(df.a))

self.assertEqual(df.filter('b = "x"').collect(), [Row(a=1, b='x')])

def test_udf_without_arguments(self):

self.spark.catalog.registerFunction("foo", lambda: "bar")

[row] = self.spark.sql("SELECT foo()").collect()

self.assertEqual(row[0], "bar")

def test_udf_with_array_type(self):

d = [Row(l=list(range(3)), d={"key": list(range(5))})]

rdd = self.sc.parallelize(d)

self.spark.createDataFrame(rdd).createOrReplaceTempView("test")

self.spark.catalog.registerFunction("copylist", lambda l: list(l), ArrayType(IntegerType()))

self.spark.catalog.registerFunction("maplen", lambda d: len(d), IntegerType())

[(l1, l2)] = self.spark.sql("select copylist(l), maplen(d) from test").collect()

self.assertEqual(list(range(3)), l1)

self.assertEqual(1, l2)

def test_broadcast_in_udf(self):

bar = {"a": "aa", "b": "bb", "c": "abc"}

foo = self.sc.broadcast(bar)

self.spark.catalog.registerFunction("MYUDF", lambda x: foo.value[x] if x else '')

[res] = self.spark.sql("SELECT MYUDF('c')").collect()

self.assertEqual("abc", res[0])

[res] = self.spark.sql("SELECT MYUDF('')").collect()

self.assertEqual("", res[0])

def test_udf_with_aggregate_function(self):

df = self.spark.createDataFrame([(1, "1"), (2, "2"), (1, "2"), (1, "2")], ["key", "value"])

from pyspark.sql.functions import udf, col, sum

from pyspark.sql.types import BooleanType

my_filter = udf(lambda a: a == 1, BooleanType())

sel = df.select(col("key")).distinct().filter(my_filter(col("key")))

self.assertEqual(sel.collect(), [Row(key=1)])

my_copy = udf(lambda x: x, IntegerType())

my_add = udf(lambda a, b: int(a + b), IntegerType())

my_strlen = udf(lambda x: len(x), IntegerType())

sel = df.groupBy(my_copy(col("key")).alias("k"))\

.agg(sum(my_strlen(col("value"))).alias("s"))\

.select(my_add(col("k"), col("s")).alias("t"))

self.assertEqual(sel.collect(), [Row(t=4), Row(t=3)])

def test_udf_in_generate(self):

from pyspark.sql.functions import udf, explode

df = self.spark.range(5)

f = udf(lambda x: list(range(x)), ArrayType(LongType()))

row = df.select(explode(f(*df))).groupBy().sum().first()

self.assertEqual(row[0], 10)

def test_udf_with_order_by_and_limit(self):

from pyspark.sql.functions import udf

my_copy = udf(lambda x: x, IntegerType())

df = self.spark.range(10).orderBy("id")

res = df.select(df.id, my_copy(df.id).alias("copy")).limit(1)

res.explain(True)

self.assertEqual(res.collect(), [Row(id=0, copy=0)])

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())

df.createOrReplaceTempView("temp")

df = self.spark.sql("select foo from temp")

df.count()

df.collect()

def test_apply_schema_to_row(self):

df = self.spark.read.json(self.sc.parallelize(["""{"a":2}"""]))

df2 = self.spark.createDataFrame(df.rdd.map(lambda x: x), df.schema)

self.assertEqual(df.collect(), df2.collect())

rdd = self.sc.parallelize(range(10)).map(lambda x: Row(a=x))

df3 = self.spark.createDataFrame(rdd, df.schema)

self.assertEqual(10, df3.count())

def test_infer_schema_to_local(self):

input = [{"a": 1}, {"b": "coffee"}]

rdd = self.sc.parallelize(input)

df = self.spark.createDataFrame(input)

df2 = self.spark.createDataFrame(rdd, samplingRatio=1.0)

self.assertEqual(df.schema, df2.schema)

rdd = self.sc.parallelize(range(10)).map(lambda x: Row(a=x, b=None))

df3 = self.spark.createDataFrame(rdd, df.schema)

self.assertEqual(10, df3.count())

def test_apply_schema_to_dict_and_rows(self):

schema = StructType().add("b", StringType()).add("a", IntegerType())

input = [{"a": 1}, {"b": "coffee"}]

rdd = self.sc.parallelize(input)

for verify in [False, True]:

df = self.spark.createDataFrame(input, schema, verifySchema=verify)

df2 = self.spark.createDataFrame(rdd, schema, verifySchema=verify)

self.assertEqual(df.schema, df2.schema)

rdd = self.sc.parallelize(range(10)).map(lambda x: Row(a=x, b=None))

df3 = self.spark.createDataFrame(rdd, schema, verifySchema=verify)

self.assertEqual(10, df3.count())

input = [Row(a=x, b=str(x)) for x in range(10)]

df4 = self.spark.createDataFrame(input, schema, verifySchema=verify)

self.assertEqual(10, df4.count())

def test_create_dataframe_schema_mismatch(self):

input = [Row(a=1)]

rdd = self.sc.parallelize(range(3)).map(lambda i: Row(a=i))

schema = StructType([StructField("a", IntegerType()), StructField("b", StringType())])

df = self.spark.createDataFrame(rdd, schema)

self.assertRaises(Exception, lambda: df.show())

def test_serialize_nested_array_and_map(self):

d = [Row(l=[Row(a=1, b='s')], d={"key": Row(c=1.0, d="2")})]

rdd = self.sc.parallelize(d)

df = self.spark.createDataFrame(rdd)

row = df.head()

self.assertEqual(1, len(row.l))

self.assertEqual(1, row.l[0].a)

self.assertEqual("2", row.d["key"].d)

l = df.rdd.map(lambda x: x.l).first()

self.assertEqual(1, len(l))

self.assertEqual('s', l[0].b)

d = df.rdd.map(lambda x: x.d).first()

self.assertEqual(1, len(d))

self.assertEqual(1.0, d["key"].c)

row = df.rdd.map(lambda x: x.d["key"]).first()

self.assertEqual(1.0, row.c)

self.assertEqual("2", row.d)

def test_infer_schema(self):

d = [Row(l=[], d={}, s=None),

Row(l=[Row(a=1, b='s')], d={"key": Row(c=1.0, d="2")}, s="")]

rdd = self.sc.parallelize(d)

df = self.spark.createDataFrame(rdd)

self.assertEqual([], df.rdd.map(lambda r: r.l).first())

self.assertEqual([None, ""], df.rdd.map(lambda r: r.s).collect())

df.createOrReplaceTempView("test")

result = self.spark.sql("SELECT l[0].a from test where d['key'].d = '2'")

self.assertEqual(1, result.head()[0])

df2 = self.spark.createDataFrame(rdd, samplingRatio=1.0)

self.assertEqual(df.schema, df2.schema)

self.assertEqual({}, df2.rdd.map(lambda r: r.d).first())

self.assertEqual([None, ""], df2.rdd.map(lambda r: r.s).collect())

df2.createOrReplaceTempView("test2")

result = self.spark.sql("SELECT l[0].a from test2 where d['key'].d = '2'")

self.assertEqual(1, result.head()[0])

def test_infer_nested_schema(self):

NestedRow = Row("f1", "f2")

nestedRdd1 = self.sc.parallelize([NestedRow([1, 2], {"row1": 1.0}),

NestedRow([2, 3], {"row2": 2.0})])

df = self.spark.createDataFrame(nestedRdd1)

self.assertEqual(Row(f1=[1, 2], f2={u'row1': 1.0}), df.collect()[0])

nestedRdd2 = self.sc.parallelize([NestedRow([[1, 2], [2, 3]], [1, 2]),

NestedRow([[2, 3], [3, 4]], [2, 3])])

df = self.spark.createDataFrame(nestedRdd2)

self.assertEqual(Row(f1=[[1, 2], [2, 3]], f2=[1, 2]), df.collect()[0])

from collections import namedtuple

CustomRow = namedtuple('CustomRow', 'field1 field2')

rdd = self.sc.parallelize([CustomRow(field1=1, field2="row1"),

CustomRow(field1=2, field2="row2"),

CustomRow(field1=3, field2="row3")])

df = self.spark.createDataFrame(rdd)

self.assertEqual(Row(field1=1, field2=u'row1'), df.first())

def test_create_dataframe_from_objects(self):

data = [MyObject(1, "1"), MyObject(2, "2")]

df = self.spark.createDataFrame(data)

self.assertEqual(df.dtypes, [("key", "bigint"), ("value", "string")])

self.assertEqual(df.first(), Row(key=1, value="1"))

def test_select_null_literal(self):

df = self.spark.sql("select null as col")

self.assertEqual(Row(col=None), df.first())

def test_apply_schema(self):

from datetime import date, datetime

rdd = self.sc.parallelize([(127, -128, -32768, 32767, 2147483647, 1.0,

date(2010, 1, 1), datetime(2010, 1, 1, 1, 1, 1),

{"a": 1}, (2,), [1, 2, 3], None)])

schema = StructType([

StructField("byte1", ByteType(), False),

StructField("byte2", ByteType(), False),

StructField("short1", ShortType(), False),

StructField("short2", ShortType(), False),

StructField("int1", IntegerType(), False),

StructField("float1", FloatType(), False),

StructField("date1", DateType(), False),

StructField("time1", TimestampType(), False),

StructField("map1", MapType(StringType(), IntegerType(), False), False),

StructField("struct1", StructType([StructField("b", ShortType(), False)]), False),

StructField("list1", ArrayType(ByteType(), False), False),

StructField("null1", DoubleType(), True)])

df = self.spark.createDataFrame(rdd, schema)

results = df.rdd.map(lambda x: (x.byte1, x.byte2, x.short1, x.short2, x.int1, x.float1,

x.date1, x.time1, x.map1["a"], x.struct1.b, x.list1, x.null1))

r = (127, -128, -32768, 32767, 2147483647, 1.0, date(2010, 1, 1),

datetime(2010, 1, 1, 1, 1, 1), 1, 2, [1, 2, 3], None)

self.assertEqual(r, results.first())

df.createOrReplaceTempView("table2")

r = self.spark.sql("SELECT byte1 - 1 AS byte1, byte2 + 1 AS byte2, " +

"short1 + 1 AS short1, short2 - 1 AS short2, int1 - 1 AS int1, " +

"float1 + 1.5 as float1 FROM table2").first()

self.assertEqual((126, -127, -32767, 32766, 2147483646, 2.5), tuple(r))

from pyspark.sql.types import _parse_schema_abstract, _infer_schema_type

rdd = self.sc.parallelize([(127, -32768, 1.0, datetime(2010, 1, 1, 1, 1, 1),

{"a": 1}, (2,), [1, 2, 3])])

abstract = "byte1 short1 float1 time1 map1{} struct1(b) list1[]"

schema = _parse_schema_abstract(abstract)

typedSchema = _infer_schema_type(rdd.first(), schema)

df = self.spark.createDataFrame(rdd, typedSchema)

r = (127, -32768, 1.0, datetime(2010, 1, 1, 1, 1, 1), {"a": 1}, Row(b=2), [1, 2, 3])

self.assertEqual(r, tuple(df.first()))

def test_struct_in_map(self):

d = [Row(m={Row(i=1): Row(s="")})]

df = self.sc.parallelize(d).toDF()

k, v = list(df.head().m.items())[0]

self.assertEqual(1, k.i)

self.assertEqual("", v.s)

def test_convert_row_to_dict(self):

row = Row(l=[Row(a=1, b='s')], d={"key": Row(c=1.0, d="2")})

self.assertEqual(1, row.asDict()['l'][0].a)

df = self.sc.parallelize([row]).toDF()

df.createOrReplaceTempView("test")

row = self.spark.sql("select l, d from test").head()

self.assertEqual(1, row.asDict()["l"][0].a)

self.assertEqual(1.0, row.asDict()['d']['key'].c)

def test_udt(self):

from pyspark.sql.types import _parse_datatype_json_string, _infer_type, _verify_type

from pyspark.sql.tests import ExamplePointUDT, ExamplePoint

def check_datatype(datatype):

pickled = pickle.loads(pickle.dumps(datatype))

assert datatype == pickled

scala_datatype = self.spark._jsparkSession.parseDataType(datatype.json())

python_datatype = _parse_datatype_json_string(scala_datatype.json())

assert datatype == python_datatype

check_datatype(ExamplePointUDT())

structtype_with_udt = StructType([StructField("label", DoubleType(), False),

StructField("point", ExamplePointUDT(), False)])

check_datatype(structtype_with_udt)

p = ExamplePoint(1.0, 2.0)

self.assertEqual(_infer_type(p), ExamplePointUDT())

_verify_type(ExamplePoint(1.0, 2.0), ExamplePointUDT())

self.assertRaises(ValueError, lambda: _verify_type([1.0, 2.0], ExamplePointUDT()))

check_datatype(PythonOnlyUDT())

structtype_with_udt = StructType([StructField("label", DoubleType(), False),

StructField("point", PythonOnlyUDT(), False)])

check_datatype(structtype_with_udt)

p = PythonOnlyPoint(1.0, 2.0)

self.assertEqual(_infer_type(p), PythonOnlyUDT())

_verify_type(PythonOnlyPoint(1.0, 2.0), PythonOnlyUDT())

self.assertRaises(ValueError, lambda: _verify_type([1.0, 2.0], PythonOnlyUDT()))

def test_simple_udt_in_df(self):

schema = StructType().add("key", LongType()).add("val", PythonOnlyUDT())

df = self.spark.createDataFrame(

[(i % 3, PythonOnlyPoint(float(i), float(i))) for i in range(10)],

schema=schema)

df.show()

def test_nested_udt_in_df(self):

schema = StructType().add("key", LongType()).add("val", ArrayType(PythonOnlyUDT()))

df = self.spark.createDataFrame(

[(i % 3, [PythonOnlyPoint(float(i), float(i))]) for i in range(10)],

schema=schema)

df.collect()

schema = StructType().add("key", LongType()).add("val",

MapType(LongType(), PythonOnlyUDT()))

df = self.spark.createDataFrame(

[(i % 3, {i % 3: PythonOnlyPoint(float(i + 1), float(i + 1))}) for i in range(10)],

schema=schema)

df.collect()

def test_complex_nested_udt_in_df(self):

from pyspark.sql.functions import udf

schema = StructType().add("key", LongType()).add("val", PythonOnlyUDT())

df = self.spark.createDataFrame(

[(i % 3, PythonOnlyPoint(float(i), float(i))) for i in range(10)],

schema=schema)

df.collect()

gd = df.groupby("key").agg({"val": "collect_list"})

gd.collect()

udf = udf(lambda k, v: [(k, v[0])], ArrayType(df.schema))

gd.select(udf(*gd)).collect()

def test_udt_with_none(self):

df = self.spark.range(0, 10, 1, 1)

def myudf(x):

if x > 0:

return PythonOnlyPoint(float(x), float(x))

self.spark.catalog.registerFunction("udf", myudf, PythonOnlyUDT())

rows = [r[0] for r in df.selectExpr("udf(id)").take(2)]

self.assertEqual(rows, [None, PythonOnlyPoint(1, 1)])

def test_infer_schema_with_udt(self):

from pyspark.sql.tests import ExamplePoint, ExamplePointUDT

row = Row(label=1.0, point=ExamplePoint(1.0, 2.0))

df = self.spark.createDataFrame([row])

schema = df.schema

field = [f for f in schema.fields if f.name == "point"][0]

self.assertEqual(type(field.dataType), ExamplePointUDT)

df.createOrReplaceTempView("labeled_point")

point = self.spark.sql("SELECT point FROM labeled_point").head().point

self.assertEqual(point, ExamplePoint(1.0, 2.0))

row = Row(label=1.0, point=PythonOnlyPoint(1.0, 2.0))

df = self.spark.createDataFrame([row])

schema = df.schema

field = [f for f in schema.fields if f.name == "point"][0]

self.assertEqual(type(field.dataType), PythonOnlyUDT)

df.createOrReplaceTempView("labeled_point")

point = self.spark.sql("SELECT point FROM labeled_point").head().point

self.assertEqual(point, PythonOnlyPoint(1.0, 2.0))

def test_apply_schema_with_udt(self):

from pyspark.sql.tests import ExamplePoint, ExamplePointUDT

row = (1.0, ExamplePoint(1.0, 2.0))

schema = StructType([StructField("label", DoubleType(), False),

StructField("point", ExamplePointUDT(), False)])

df = self.spark.createDataFrame([row], schema)

point = df.head().point

self.assertEqual(point, ExamplePoint(1.0, 2.0))

row = (1.0, PythonOnlyPoint(1.0, 2.0))

schema = StructType([StructField("label", DoubleType(), False),

StructField("point", PythonOnlyUDT(), False)])

df = self.spark.createDataFrame([row], schema)

point = df.head().point

self.assertEqual(point, PythonOnlyPoint(1.0, 2.0))

def test_udf_with_udt(self):

from pyspark.sql.tests import ExamplePoint, ExamplePointUDT

row = Row(label=1.0, point=ExamplePoint(1.0, 2.0))

df = self.spark.createDataFrame([row])

self.assertEqual(1.0, df.rdd.map(lambda r: r.point.x).first())

udf = UserDefinedFunction(lambda p: p.y, DoubleType())

self.assertEqual(2.0, df.select(udf(df.point)).first()[0])

udf2 = UserDefinedFunction(lambda p: ExamplePoint(p.x + 1, p.y + 1), ExamplePointUDT())

self.assertEqual(ExamplePoint(2.0, 3.0), df.select(udf2(df.point)).first()[0])

row = Row(label=1.0, point=PythonOnlyPoint(1.0, 2.0))

df = self.spark.createDataFrame([row])

self.assertEqual(1.0, df.rdd.map(lambda r: r.point.x).first())

udf = UserDefinedFunction(lambda p: p.y, DoubleType())

self.assertEqual(2.0, df.select(udf(df.point)).first()[0])

udf2 = UserDefinedFunction(lambda p: PythonOnlyPoint(p.x + 1, p.y + 1), PythonOnlyUDT())

self.assertEqual(PythonOnlyPoint(2.0, 3.0), df.select(udf2(df.point)).first()[0])

def test_parquet_with_udt(self):

from pyspark.sql.tests import ExamplePoint, ExamplePointUDT

row = Row(label=1.0, point=ExamplePoint(1.0, 2.0))

df0 = self.spark.createDataFrame([row])

output_dir = os.path.join(self.tempdir.name, "labeled_point")

df0.write.parquet(output_dir)

df1 = self.spark.read.parquet(output_dir)

point = df1.head().point

self.assertEqual(point, ExamplePoint(1.0, 2.0))

row = Row(label=1.0, point=PythonOnlyPoint(1.0, 2.0))

df0 = self.spark.createDataFrame([row])

df0.write.parquet(output_dir, mode='overwrite')

df1 = self.spark.read.parquet(output_dir)

point = df1.head().point

self.assertEqual(point, PythonOnlyPoint(1.0, 2.0))

def test_union_with_udt(self):

from pyspark.sql.tests import ExamplePoint, ExamplePointUDT

row1 = (1.0, ExamplePoint(1.0, 2.0))

row2 = (2.0, ExamplePoint(3.0, 4.0))

schema = StructType([StructField("label", DoubleType(), False),

StructField("point", ExamplePointUDT(), False)])

df1 = self.spark.createDataFrame([row1], schema)

df2 = self.spark.createDataFrame([row2], schema)

result = df1.union(df2).orderBy("label").collect()

self.assertEqual(

result,

[

Row(label=1.0, point=ExamplePoint(1.0, 2.0)),

Row(label=2.0, point=ExamplePoint(3.0, 4.0))

]

)

def test_column_operators(self):

ci = self.df.key

cs = self.df.value

c = ci == cs

self.assertTrue(isinstance((- ci - 1 - 2) % 3 * 2.5 / 3.5, Column))

rcc = (1 + ci), (1 - ci), (1 * ci), (1 / ci), (1 % ci), (1 ** ci), (ci ** 1)

self.assertTrue(all(isinstance(c, Column) for c in rcc))

cb = [ci == 5, ci != 0, ci > 3, ci < 4, ci >= 0, ci <= 7]

self.assertTrue(all(isinstance(c, Column) for c in cb))

cbool = (ci & ci), (ci | ci), (~ci)

self.assertTrue(all(isinstance(c, Column) for c in cbool))

css = cs.like('a'), cs.rlike('a'), cs.asc(), cs.desc(), cs.startswith('a'), cs.endswith('a')

self.assertTrue(all(isinstance(c, Column) for c in css))

self.assertTrue(isinstance(ci.cast(LongType()), Column))

def test_column_select(self):

df = self.df

self.assertEqual(self.testData, df.select("*").collect())

self.assertEqual(self.testData, df.select(df.key, df.value).collect())

self.assertEqual([Row(value='1')], df.where(df.key == 1).select(df.value).collect())

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_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.approxCountDistinct(df.key)).first()[0])

self.assertEqual(100, g.agg(functions.countDistinct(df.value)).first()[0])

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) for i in range(10)]).toDF()

aq = df.stat.approxQuantile("a", [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))

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_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

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_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_struct_type(self):

from pyspark.sql.types import StructType, StringType, StructField

struct1 = StructType().add("f1", StringType(), True).add("f2", StringType(), True, None)

struct2 = StructType([StructField("f1", StringType(), True),

StructField("f2", StringType(), True, None)])

self.assertEqual(struct1, struct2)

struct1 = StructType().add("f1", StringType(), True).add("f2", StringType(), True, None)

struct2 = StructType([StructField("f1", StringType(), True)])

self.assertNotEqual(struct1, struct2)

struct1 = (StructType().add(StructField("f1", StringType(), True))

.add(StructField("f2", StringType(), True, None)))

struct2 = StructType([StructField("f1", StringType(), True),

StructField("f2", StringType(), True, None)])

self.assertEqual(struct1, struct2)

struct1 = (StructType().add(StructField("f1", StringType(), True))

.add(StructField("f2", StringType(), True, None)))

struct2 = StructType([StructField("f1", StringType(), True)])

self.assertNotEqual(struct1, struct2)

# Catch exception raised during improper construction

with self.assertRaises(ValueError):

struct1 = StructType().add("name")

struct1 = StructType().add("f1", StringType(), True).add("f2", StringType(), True, None)

for field in struct1:

self.assertIsInstance(field, StructField)

struct1 = StructType().add("f1", StringType(), True).add("f2", StringType(), True, None)

self.assertEqual(len(struct1), 2)

struct1 = StructType().add("f1", StringType(), True).add("f2", StringType(), True, None)

self.assertIs(struct1["f1"], struct1.fields[0])

self.assertIs(struct1[0], struct1.fields[0])

self.assertEqual(struct1[0:1], StructType(struct1.fields[0:1]))

with self.assertRaises(KeyError):

not_a_field = struct1["f9"]

with self.assertRaises(IndexError):

not_a_field = struct1[9]

with self.assertRaises(TypeError):

not_a_field = struct1[9.9]

def test_metadata_null(self):

from pyspark.sql.types import StructType, StringType, StructField

schema = StructType([StructField("f1", StringType(), True, None),

StructField("f2", StringType(), True, {'a': None})])

rdd = self.sc.parallelize([["a", "b"], ["c", "d"]])

self.spark.createDataFrame(rdd, schema)

def test_save_and_load(self):

df = self.df

tmpPath = tempfile.mkdtemp()

shutil.rmtree(tmpPath)

df.write.json(tmpPath)

actual = self.spark.read.json(tmpPath)

self.assertEqual(sorted(df.collect()), sorted(actual.collect()))

schema = StructType([StructField("value", StringType(), True)])

actual = self.spark.read.json(tmpPath, schema)

self.assertEqual(sorted(df.select("value").collect()), sorted(actual.collect()))

df.write.json(tmpPath, "overwrite")

actual = self.spark.read.json(tmpPath)

self.assertEqual(sorted(df.collect()), sorted(actual.collect()))

df.write.save(format="json", mode="overwrite", path=tmpPath,

noUse="this options will not be used in save.")

actual = self.spark.read.load(format="json", path=tmpPath,

noUse="this options will not be used in load.")

self.assertEqual(sorted(df.collect()), sorted(actual.collect()))

defaultDataSourceName = self.spark.conf.get("spark.sql.sources.default",

"org.apache.spark.sql.parquet")

self.spark.sql("SET spark.sql.sources.default=org.apache.spark.sql.json")

actual = self.spark.read.load(path=tmpPath)

self.assertEqual(sorted(df.collect()), sorted(actual.collect()))

self.spark.sql("SET spark.sql.sources.default=" + defaultDataSourceName)

csvpath = os.path.join(tempfile.mkdtemp(), 'data')

df.write.option('quote', None).format('csv').save(csvpath)

shutil.rmtree(tmpPath)

def test_save_and_load_builder(self):

df = self.df

tmpPath = tempfile.mkdtemp()

shutil.rmtree(tmpPath)

df.write.json(tmpPath)

actual = self.spark.read.json(tmpPath)

self.assertEqual(sorted(df.collect()), sorted(actual.collect()))

schema = StructType([StructField("value", StringType(), True)])

actual = self.spark.read.json(tmpPath, schema)

self.assertEqual(sorted(df.select("value").collect()), sorted(actual.collect()))

df.write.mode("overwrite").json(tmpPath)

actual = self.spark.read.json(tmpPath)

self.assertEqual(sorted(df.collect()), sorted(actual.collect()))

df.write.mode("overwrite").options(noUse="this options will not be used in save.")\

.option("noUse", "this option will not be used in save.")\

.format("json").save(path=tmpPath)

actual =\

self.spark.read.format("json")\

.load(path=tmpPath, noUse="this options will not be used in load.")

self.assertEqual(sorted(df.collect()), sorted(actual.collect()))

defaultDataSourceName = self.spark.conf.get("spark.sql.sources.default",

"org.apache.spark.sql.parquet")

self.spark.sql("SET spark.sql.sources.default=org.apache.spark.sql.json")

actual = self.spark.read.load(path=tmpPath)

self.assertEqual(sorted(df.collect()), sorted(actual.collect()))

self.spark.sql("SET spark.sql.sources.default=" + defaultDataSourceName)

shutil.rmtree(tmpPath)