| # |
| # 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 array |
| import datetime |
| import unittest |
| import random |
| import string |
| |
| from pyspark.errors import PySparkValueError |
| from pyspark.sql.types import ( |
| StructType, |
| StructField, |
| StringType, |
| IntegerType, |
| LongType, |
| MapType, |
| ArrayType, |
| Row, |
| ) |
| from pyspark.testing.objects import MyObject, PythonOnlyUDT |
| from pyspark.testing.connectutils import should_test_connect, ReusedMixedTestCase |
| from pyspark.testing.pandasutils import PandasOnSparkTestUtils |
| |
| if should_test_connect: |
| import pandas as pd |
| import numpy as np |
| from pyspark.sql import functions as SF |
| from pyspark.sql.connect import functions as CF |
| from pyspark.errors.exceptions.connect import ParseException |
| |
| |
| class SparkConnectCreationTests(ReusedMixedTestCase, PandasOnSparkTestUtils): |
| def test_with_local_data(self): |
| """SPARK-41114: Test creating a dataframe using local data""" |
| pdf = pd.DataFrame({"a": [1, 2, 3], "b": ["a", "b", "c"]}) |
| df = self.connect.createDataFrame(pdf) |
| rows = df.filter(df.a == CF.lit(3)).collect() |
| self.assertTrue(len(rows) == 1) |
| self.assertEqual(rows[0][0], 3) |
| self.assertEqual(rows[0][1], "c") |
| |
| # Check correct behavior for empty DataFrame |
| pdf = pd.DataFrame({"a": []}) |
| with self.assertRaises(ValueError): |
| self.connect.createDataFrame(pdf) |
| |
| def test_with_local_ndarray(self): |
| """SPARK-41446: Test creating a dataframe using local list""" |
| data = np.array([[1, 2, 3, 4], [5, 6, 7, 8]]) |
| |
| sdf = self.spark.createDataFrame(data) |
| cdf = self.connect.createDataFrame(data) |
| self.assertEqual(sdf.schema, cdf.schema) |
| self.assert_eq(sdf.toPandas(), cdf.toPandas()) |
| |
| for schema in [ |
| StructType( |
| [ |
| StructField("col1", IntegerType(), True), |
| StructField("col2", IntegerType(), True), |
| StructField("col3", IntegerType(), True), |
| StructField("col4", IntegerType(), True), |
| ] |
| ), |
| "struct<col1 int, col2 int, col3 int, col4 int>", |
| "col1 int, col2 int, col3 int, col4 int", |
| "col1 int, col2 long, col3 string, col4 long", |
| "col1 int, col2 string, col3 short, col4 long", |
| ["a", "b", "c", "d"], |
| ("x1", "x2", "x3", "x4"), |
| ]: |
| with self.subTest(schema=schema): |
| sdf = self.spark.createDataFrame(data, schema=schema) |
| cdf = self.connect.createDataFrame(data, schema=schema) |
| |
| self.assertEqual(sdf.schema, cdf.schema) |
| self.assert_eq(sdf.toPandas(), cdf.toPandas()) |
| |
| with self.assertRaises(PySparkValueError) as pe: |
| self.connect.createDataFrame(data, ["a", "b", "c", "d", "e"]) |
| |
| self.check_error( |
| exception=pe.exception, |
| errorClass="AXIS_LENGTH_MISMATCH", |
| messageParameters={"expected_length": "5", "actual_length": "4"}, |
| ) |
| |
| with self.assertRaises(ParseException): |
| self.connect.createDataFrame(data, "col1 magic_type, col2 int, col3 int, col4 int") |
| |
| with self.assertRaises(PySparkValueError) as pe: |
| self.connect.createDataFrame(data, "col1 int, col2 int, col3 int") |
| |
| self.check_error( |
| exception=pe.exception, |
| errorClass="AXIS_LENGTH_MISMATCH", |
| messageParameters={"expected_length": "3", "actual_length": "4"}, |
| ) |
| |
| # test 1 dim ndarray |
| data = np.array([1.0, 2.0, np.nan, 3.0, 4.0, float("NaN"), 5.0]) |
| self.assertEqual(data.ndim, 1) |
| |
| sdf = self.spark.createDataFrame(data) |
| cdf = self.connect.createDataFrame(data) |
| self.assertEqual(sdf.schema, cdf.schema) |
| self.assert_eq(sdf.toPandas(), cdf.toPandas()) |
| |
| def test_with_local_list(self): |
| """SPARK-41446: Test creating a dataframe using local list""" |
| data = [[1, 2, 3, 4]] |
| |
| sdf = self.spark.createDataFrame(data) |
| cdf = self.connect.createDataFrame(data) |
| self.assertEqual(sdf.schema, cdf.schema) |
| self.assert_eq(sdf.toPandas(), cdf.toPandas()) |
| |
| for schema in [ |
| "struct<col1 int, col2 int, col3 int, col4 int>", |
| "col1 int, col2 int, col3 int, col4 int", |
| "col1 int, col2 long, col3 string, col4 long", |
| "col1 int, col2 string, col3 short, col4 long", |
| ["a", "b", "c", "d"], |
| ("x1", "x2", "x3", "x4"), |
| ]: |
| sdf = self.spark.createDataFrame(data, schema=schema) |
| cdf = self.connect.createDataFrame(data, schema=schema) |
| |
| self.assertEqual(sdf.schema, cdf.schema) |
| self.assert_eq(sdf.toPandas(), cdf.toPandas()) |
| |
| with self.assertRaises(PySparkValueError) as pe: |
| self.connect.createDataFrame(data, ["a", "b", "c", "d", "e"]) |
| |
| self.check_error( |
| exception=pe.exception, |
| errorClass="AXIS_LENGTH_MISMATCH", |
| messageParameters={"expected_length": "5", "actual_length": "4"}, |
| ) |
| |
| with self.assertRaises(ParseException): |
| self.connect.createDataFrame(data, "col1 magic_type, col2 int, col3 int, col4 int") |
| |
| with self.assertRaises(PySparkValueError) as pe: |
| self.connect.createDataFrame(data, "col1 int, col2 int, col3 int") |
| |
| self.check_error( |
| exception=pe.exception, |
| errorClass="AXIS_LENGTH_MISMATCH", |
| messageParameters={"expected_length": "3", "actual_length": "4"}, |
| ) |
| |
| def test_with_local_rows(self): |
| # SPARK-41789, SPARK-41810: Test creating a dataframe with list of rows and dictionaries |
| rows = [ |
| Row(course="dotNET", year=2012, earnings=10000), |
| Row(course="Java", year=2012, earnings=20000), |
| Row(course="dotNET", year=2012, earnings=5000), |
| Row(course="dotNET", year=2013, earnings=48000), |
| Row(course="Java", year=2013, earnings=30000), |
| Row(course="Scala", year=2022, earnings=None), |
| ] |
| dicts = [row.asDict() for row in rows] |
| |
| for data in [rows, dicts]: |
| sdf = self.spark.createDataFrame(data) |
| cdf = self.connect.createDataFrame(data) |
| |
| self.assertEqual(sdf.schema, cdf.schema) |
| self.assert_eq(sdf.toPandas(), cdf.toPandas()) |
| |
| # test with rename |
| sdf = self.spark.createDataFrame(data, schema=["a", "b", "c"]) |
| cdf = self.connect.createDataFrame(data, schema=["a", "b", "c"]) |
| |
| self.assertEqual(sdf.schema, cdf.schema) |
| self.assert_eq(sdf.toPandas(), cdf.toPandas()) |
| |
| def test_streaming_local_relation(self): |
| threshold_conf = "spark.sql.session.localRelationCacheThreshold" |
| old_threshold = self.connect.conf.get(threshold_conf) |
| threshold = 1024 * 1024 |
| self.connect.conf.set(threshold_conf, threshold) |
| try: |
| suffix = "abcdef" |
| letters = string.ascii_lowercase |
| str = "".join(random.choice(letters) for i in range(threshold)) + suffix |
| data = [[0, str], [1, str]] |
| for i in range(0, 2): |
| cdf = self.connect.createDataFrame(data, ["a", "b"]) |
| self.assert_eq(cdf.count(), len(data)) |
| self.assert_eq(cdf.filter(f"endsWith(b, '{suffix}')").isEmpty(), False) |
| finally: |
| self.connect.conf.set(threshold_conf, old_threshold) |
| |
| def test_with_atom_type(self): |
| for data in [[(1), (2), (3)], [1, 2, 3]]: |
| for schema in ["long", "int", "short"]: |
| sdf = self.spark.createDataFrame(data, schema=schema) |
| cdf = self.connect.createDataFrame(data, schema=schema) |
| |
| self.assertEqual(sdf.schema, cdf.schema) |
| self.assert_eq(sdf.toPandas(), cdf.toPandas()) |
| |
| def test_with_none_and_nan(self): |
| # SPARK-41855: make createDataFrame support None and NaN |
| # SPARK-41814: test with eqNullSafe |
| data1 = [Row(id=1, value=float("NaN")), Row(id=2, value=42.0), Row(id=3, value=None)] |
| data2 = [Row(id=1, value=np.nan), Row(id=2, value=42.0), Row(id=3, value=None)] |
| data3 = [ |
| {"id": 1, "value": float("NaN")}, |
| {"id": 2, "value": 42.0}, |
| {"id": 3, "value": None}, |
| ] |
| data4 = [{"id": 1, "value": np.nan}, {"id": 2, "value": 42.0}, {"id": 3, "value": None}] |
| data5 = [(1, float("NaN")), (2, 42.0), (3, None)] |
| data6 = [(1, np.nan), (2, 42.0), (3, None)] |
| data7 = np.array([[1, float("NaN")], [2, 42.0], [3, None]]) |
| data8 = np.array([[1, np.nan], [2, 42.0], [3, None]]) |
| |
| # +---+-----+ |
| # | id|value| |
| # +---+-----+ |
| # | 1| NaN| |
| # | 2| 42.0| |
| # | 3| NULL| |
| # +---+-----+ |
| |
| for data in [data1, data2, data3, data4, data5, data6, data7, data8]: |
| if isinstance(data[0], (Row, dict)): |
| # data1, data2, data3, data4 |
| cdf = self.connect.createDataFrame(data) |
| sdf = self.spark.createDataFrame(data) |
| else: |
| # data5, data6, data7, data8 |
| cdf = self.connect.createDataFrame(data, schema=["id", "value"]) |
| sdf = self.spark.createDataFrame(data, schema=["id", "value"]) |
| |
| self.assert_eq(cdf.toPandas(), sdf.toPandas()) |
| |
| self.assert_eq( |
| cdf.select( |
| cdf["value"].eqNullSafe(None), |
| cdf["value"].eqNullSafe(float("NaN")), |
| cdf["value"].eqNullSafe(42.0), |
| ).toPandas(), |
| sdf.select( |
| sdf["value"].eqNullSafe(None), |
| sdf["value"].eqNullSafe(float("NaN")), |
| sdf["value"].eqNullSafe(42.0), |
| ).toPandas(), |
| ) |
| |
| # SPARK-41851: test with nanvl |
| data = [(1.0, float("nan")), (float("nan"), 2.0)] |
| |
| cdf = self.connect.createDataFrame(data, ("a", "b")) |
| sdf = self.spark.createDataFrame(data, ("a", "b")) |
| |
| self.assert_eq(cdf.toPandas(), sdf.toPandas()) |
| |
| self.assert_eq( |
| cdf.select( |
| CF.nanvl("a", "b").alias("r1"), CF.nanvl(cdf.a, cdf.b).alias("r2") |
| ).toPandas(), |
| sdf.select( |
| SF.nanvl("a", "b").alias("r1"), SF.nanvl(sdf.a, sdf.b).alias("r2") |
| ).toPandas(), |
| ) |
| |
| # SPARK-41852: test with pmod |
| data = [ |
| (1.0, float("nan")), |
| (float("nan"), 2.0), |
| (10.0, 3.0), |
| (float("nan"), float("nan")), |
| (-3.0, 4.0), |
| (-10.0, 3.0), |
| (-5.0, -6.0), |
| (7.0, -8.0), |
| (1.0, 2.0), |
| ] |
| |
| cdf = self.connect.createDataFrame(data, ("a", "b")) |
| sdf = self.spark.createDataFrame(data, ("a", "b")) |
| |
| self.assert_eq(cdf.toPandas(), sdf.toPandas()) |
| |
| self.assert_eq( |
| cdf.select(CF.pmod("a", "b")).toPandas(), |
| sdf.select(SF.pmod("a", "b")).toPandas(), |
| ) |
| |
| def test_cast_with_ddl(self): |
| data = [Row(date=datetime.date(2021, 12, 27), add=2)] |
| |
| cdf = self.connect.createDataFrame(data, "date date, add integer") |
| sdf = self.spark.createDataFrame(data, "date date, add integer") |
| |
| self.assertEqual(cdf.schema, sdf.schema) |
| |
| def test_create_empty_df(self): |
| for schema in [ |
| "STRING", |
| "x STRING", |
| "x STRING, y INTEGER", |
| StringType(), |
| StructType( |
| [ |
| StructField("x", StringType(), True), |
| StructField("y", IntegerType(), True), |
| ] |
| ), |
| ]: |
| cdf = self.connect.createDataFrame(data=[], schema=schema) |
| sdf = self.spark.createDataFrame(data=[], schema=schema) |
| |
| self.assert_eq(cdf.toPandas(), sdf.toPandas()) |
| |
| # check error |
| with self.assertRaises(PySparkValueError) as pe: |
| self.connect.createDataFrame(data=[]) |
| |
| self.check_error( |
| exception=pe.exception, |
| errorClass="CANNOT_INFER_EMPTY_SCHEMA", |
| messageParameters={}, |
| ) |
| |
| def test_create_dataframe_from_arrays(self): |
| # SPARK-42021: createDataFrame support array.array |
| data1 = [Row(a=1, b=array.array("i", [1, 2, 3]), c=array.array("d", [4, 5, 6]))] |
| data2 = [(array.array("d", [1, 2, 3]), 2, "3")] |
| data3 = [{"a": 1, "b": array.array("i", [1, 2, 3])}] |
| |
| for data in [data1, data2, data3]: |
| cdf = self.connect.createDataFrame(data) |
| sdf = self.spark.createDataFrame(data) |
| |
| self.assertEqual(cdf.schema, sdf.schema) |
| self.assertEqual(cdf.collect(), sdf.collect()) |
| |
| def test_timestampe_create_from_rows(self): |
| data = [(datetime.datetime(2016, 3, 11, 9, 0, 7), 1)] |
| |
| cdf = self.connect.createDataFrame(data, ["date", "val"]) |
| sdf = self.spark.createDataFrame(data, ["date", "val"]) |
| |
| self.assertEqual(cdf.schema, sdf.schema) |
| self.assertEqual(cdf.collect(), sdf.collect()) |
| |
| def test_create_dataframe_with_coercion(self): |
| data1 = [[1.33, 1], ["2.1", 1]] |
| data2 = [[True, 1], ["false", 1]] |
| |
| for data in [data1, data2]: |
| cdf = self.connect.createDataFrame(data, ["a", "b"]) |
| sdf = self.spark.createDataFrame(data, ["a", "b"]) |
| |
| self.assertEqual(cdf.schema, sdf.schema) |
| self.assertEqual(cdf.collect(), sdf.collect()) |
| |
| def test_nested_type_create_from_rows(self): |
| data1 = [Row(a=1, b=Row(c=2, d=Row(e=3, f=Row(g=4, h=Row(i=5)))))] |
| # root |
| # |-- a: long (nullable = true) |
| # |-- b: struct (nullable = true) |
| # | |-- c: long (nullable = true) |
| # | |-- d: struct (nullable = true) |
| # | | |-- e: long (nullable = true) |
| # | | |-- f: struct (nullable = true) |
| # | | | |-- g: long (nullable = true) |
| # | | | |-- h: struct (nullable = true) |
| # | | | | |-- i: long (nullable = true) |
| |
| data2 = [ |
| ( |
| 1, |
| "a", |
| Row( |
| a=1, |
| b=[1, 2, 3], |
| c={"a": "b"}, |
| d=Row(x=1, y="y", z=Row(o=1, p=2, q=Row(g=1.5))), |
| ), |
| ) |
| ] |
| # root |
| # |-- _1: long (nullable = true) |
| # |-- _2: string (nullable = true) |
| # |-- _3: struct (nullable = true) |
| # | |-- a: long (nullable = true) |
| # | |-- b: array (nullable = true) |
| # | | |-- element: long (containsNull = true) |
| # | |-- c: map (nullable = true) |
| # | | |-- key: string |
| # | | |-- value: string (valueContainsNull = true) |
| # | |-- d: struct (nullable = true) |
| # | | |-- x: long (nullable = true) |
| # | | |-- y: string (nullable = true) |
| # | | |-- z: struct (nullable = true) |
| # | | | |-- o: long (nullable = true) |
| # | | | |-- p: long (nullable = true) |
| # | | | |-- q: struct (nullable = true) |
| # | | | | |-- g: double (nullable = true) |
| |
| data3 = [ |
| Row( |
| a=1, |
| b=[1, 2, 3], |
| c={"a": "b"}, |
| d=Row(x=1, y="y", z=Row(1, 2, 3)), |
| e=list("hello connect"), |
| ) |
| ] |
| # root |
| # |-- a: long (nullable = true) |
| # |-- b: array (nullable = true) |
| # | |-- element: long (containsNull = true) |
| # |-- c: map (nullable = true) |
| # | |-- key: string |
| # | |-- value: string (valueContainsNull = true) |
| # |-- d: struct (nullable = true) |
| # | |-- x: long (nullable = true) |
| # | |-- y: string (nullable = true) |
| # | |-- z: struct (nullable = true) |
| # | | |-- _1: long (nullable = true) |
| # | | |-- _2: long (nullable = true) |
| # | | |-- _3: long (nullable = true) |
| # |-- e: array (nullable = true) |
| # | |-- element: string (containsNull = true) |
| |
| data4 = [ |
| { |
| "a": 1, |
| "b": Row(x=1, y=Row(z=2)), |
| "c": {"x": -1, "y": 2}, |
| "d": [1, 2, 3, 4, 5], |
| } |
| ] |
| # root |
| # |-- a: long (nullable = true) |
| # |-- b: struct (nullable = true) |
| # | |-- x: long (nullable = true) |
| # | |-- y: struct (nullable = true) |
| # | | |-- z: long (nullable = true) |
| # |-- c: map (nullable = true) |
| # | |-- key: string |
| # | |-- value: long (valueContainsNull = true) |
| # |-- d: array (nullable = true) |
| # | |-- element: long (containsNull = true) |
| |
| data5 = [ |
| { |
| "a": [Row(x=1, y="2"), Row(x=-1, y="-2")], |
| "b": [[1, 2, 3], [4, 5], [6]], |
| "c": {3: {4: {5: 6}}, 7: {8: {9: 0}}}, |
| } |
| ] |
| # root |
| # |-- a: array (nullable = true) |
| # | |-- element: struct (containsNull = true) |
| # | | |-- x: long (nullable = true) |
| # | | |-- y: string (nullable = true) |
| # |-- b: array (nullable = true) |
| # | |-- element: array (containsNull = true) |
| # | | |-- element: long (containsNull = true) |
| # |-- c: map (nullable = true) |
| # | |-- key: long |
| # | |-- value: map (valueContainsNull = true) |
| # | | |-- key: long |
| # | | |-- value: map (valueContainsNull = true) |
| # | | | |-- key: long |
| # | | | |-- value: long (valueContainsNull = true) |
| |
| for data in [data1, data2, data3, data4, data5]: |
| with self.subTest(data=data): |
| cdf = self.connect.createDataFrame(data) |
| sdf = self.spark.createDataFrame(data) |
| |
| self.assertEqual(cdf.schema, sdf.schema) |
| self.assertEqual(cdf.collect(), sdf.collect()) |
| |
| def test_create_df_from_objects(self): |
| data = [MyObject(1, "1"), MyObject(2, "2")] |
| |
| # +---+-----+ |
| # |key|value| |
| # +---+-----+ |
| # | 1| 1| |
| # | 2| 2| |
| # +---+-----+ |
| |
| cdf = self.connect.createDataFrame(data) |
| sdf = self.spark.createDataFrame(data) |
| |
| self.assertEqual(cdf.schema, sdf.schema) |
| self.assertEqual(cdf.collect(), sdf.collect()) |
| |
| def test_create_df_nullability(self): |
| data = [("asd", None)] |
| schema = StructType( |
| [ |
| StructField("name", StringType(), nullable=True), |
| StructField("age", IntegerType(), nullable=False), |
| ] |
| ) |
| |
| with self.assertRaises(PySparkValueError): |
| self.spark.createDataFrame(data, schema) |
| |
| def test_create_dataframe_from_pandas_with_ns_timestamp(self): |
| """Truncate the timestamps for nanoseconds.""" |
| from datetime import datetime, timezone, timedelta |
| from pandas import Timestamp |
| import pandas as pd |
| |
| # Nanoseconds are truncated to microseconds in the serializer |
| # Arrow will throw an error if precision is lost |
| # (i.e., nanoseconds cannot be represented in microseconds) |
| pdf = pd.DataFrame( |
| { |
| "naive": [datetime(2019, 1, 1, 0)], |
| "aware": [ |
| Timestamp( |
| year=2019, month=1, day=1, nanosecond=0, tz=timezone(timedelta(hours=-8)) |
| ) |
| ], |
| } |
| ) |
| |
| with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": False}): |
| self.assertEqual( |
| self.connect.createDataFrame(pdf).collect(), |
| self.spark.createDataFrame(pdf).collect(), |
| ) |
| |
| with self.sql_conf({"spark.sql.execution.arrow.pyspark.enabled": True}): |
| self.assertEqual( |
| self.connect.createDataFrame(pdf).collect(), |
| self.spark.createDataFrame(pdf).collect(), |
| ) |
| |
| def test_schema_has_nullable(self): |
| schema_false = StructType().add("id", IntegerType(), False) |
| cdf1 = self.connect.createDataFrame([[1]], schema=schema_false) |
| sdf1 = self.spark.createDataFrame([[1]], schema=schema_false) |
| self.assertEqual(cdf1.schema, sdf1.schema) |
| self.assertEqual(cdf1.collect(), sdf1.collect()) |
| |
| schema_true = StructType().add("id", IntegerType(), True) |
| cdf2 = self.connect.createDataFrame([[1]], schema=schema_true) |
| sdf2 = self.spark.createDataFrame([[1]], schema=schema_true) |
| self.assertEqual(cdf2.schema, sdf2.schema) |
| self.assertEqual(cdf2.collect(), sdf2.collect()) |
| |
| pdf1 = cdf1.toPandas() |
| cdf3 = self.connect.createDataFrame(pdf1, cdf1.schema) |
| sdf3 = self.spark.createDataFrame(pdf1, sdf1.schema) |
| self.assertEqual(cdf3.schema, sdf3.schema) |
| self.assertEqual(cdf3.collect(), sdf3.collect()) |
| |
| pdf2 = cdf2.toPandas() |
| cdf4 = self.connect.createDataFrame(pdf2, cdf2.schema) |
| sdf4 = self.spark.createDataFrame(pdf2, sdf2.schema) |
| self.assertEqual(cdf4.schema, sdf4.schema) |
| self.assertEqual(cdf4.collect(), sdf4.collect()) |
| |
| def test_array_has_nullable(self): |
| for schemas, data in [ |
| ( |
| [StructType().add("arr", ArrayType(IntegerType(), False), True)], |
| [Row([1, 2]), Row([3]), Row(None)], |
| ), |
| ( |
| [ |
| StructType().add("arr", ArrayType(IntegerType(), True), True), |
| "arr array<integer>", |
| ], |
| [Row([1, None]), Row([3]), Row(None)], |
| ), |
| ( |
| [StructType().add("arr", ArrayType(IntegerType(), False), False)], |
| [Row([1, 2]), Row([3])], |
| ), |
| ( |
| [ |
| StructType().add("arr", ArrayType(IntegerType(), True), False), |
| "arr array<integer> not null", |
| ], |
| [Row([1, None]), Row([3])], |
| ), |
| ]: |
| for schema in schemas: |
| with self.subTest(schema=schema): |
| cdf = self.connect.createDataFrame(data, schema=schema) |
| sdf = self.spark.createDataFrame(data, schema=schema) |
| self.assertEqual(cdf.schema, sdf.schema) |
| self.assertEqual(cdf.collect(), sdf.collect()) |
| |
| def test_map_has_nullable(self): |
| for schemas, data in [ |
| ( |
| [StructType().add("map", MapType(StringType(), IntegerType(), False), True)], |
| [Row({"a": 1, "b": 2}), Row({"a": 3}), Row(None)], |
| ), |
| ( |
| [ |
| StructType().add("map", MapType(StringType(), IntegerType(), True), True), |
| "map map<string, integer>", |
| ], |
| [Row({"a": 1, "b": None}), Row({"a": 3}), Row(None)], |
| ), |
| ( |
| [StructType().add("map", MapType(StringType(), IntegerType(), False), False)], |
| [Row({"a": 1, "b": 2}), Row({"a": 3})], |
| ), |
| ( |
| [ |
| StructType().add("map", MapType(StringType(), IntegerType(), True), False), |
| "map map<string, integer> not null", |
| ], |
| [Row({"a": 1, "b": None}), Row({"a": 3})], |
| ), |
| ]: |
| for schema in schemas: |
| with self.subTest(schema=schema): |
| cdf = self.connect.createDataFrame(data, schema=schema) |
| sdf = self.spark.createDataFrame(data, schema=schema) |
| self.assertEqual(cdf.schema, sdf.schema) |
| self.assertEqual(cdf.collect(), sdf.collect()) |
| |
| def test_struct_has_nullable(self): |
| for schemas, data in [ |
| ( |
| [ |
| StructType().add("struct", StructType().add("i", IntegerType(), False), True), |
| "struct struct<i: integer not null>", |
| ], |
| [Row(Row(1)), Row(Row(2)), Row(None)], |
| ), |
| ( |
| [ |
| StructType().add("struct", StructType().add("i", IntegerType(), True), True), |
| "struct struct<i: integer>", |
| ], |
| [Row(Row(1)), Row(Row(2)), Row(Row(None)), Row(None)], |
| ), |
| ( |
| [ |
| StructType().add("struct", StructType().add("i", IntegerType(), False), False), |
| "struct struct<i: integer not null> not null", |
| ], |
| [Row(Row(1)), Row(Row(2))], |
| ), |
| ( |
| [ |
| StructType().add("struct", StructType().add("i", IntegerType(), True), False), |
| "struct struct<i: integer> not null", |
| ], |
| [Row(Row(1)), Row(Row(2)), Row(Row(None))], |
| ), |
| ]: |
| for schema in schemas: |
| with self.subTest(schema=schema): |
| cdf = self.connect.createDataFrame(data, schema=schema) |
| sdf = self.spark.createDataFrame(data, schema=schema) |
| self.assertEqual(cdf.schema, sdf.schema) |
| self.assertEqual(cdf.collect(), sdf.collect()) |
| |
| def test_large_client_data(self): |
| # SPARK-42816 support more than 4MB message size. |
| # ~200bytes |
| cols = ["abcdefghijklmnoprstuvwxyz" for x in range(10)] |
| # 100k rows => 20MB |
| row_count = 100 * 1000 |
| rows = [cols] * row_count |
| self.assertEqual(row_count, self.connect.createDataFrame(data=rows).count()) |
| |
| def test_simple_udt(self): |
| from pyspark.ml.linalg import MatrixUDT, VectorUDT |
| |
| for schema in [ |
| StructType().add("key", LongType()).add("val", PythonOnlyUDT()), |
| StructType().add("key", LongType()).add("val", ArrayType(PythonOnlyUDT())), |
| StructType().add("key", LongType()).add("val", MapType(LongType(), PythonOnlyUDT())), |
| StructType().add("key", LongType()).add("val", PythonOnlyUDT()), |
| StructType().add("key", LongType()).add("vec", VectorUDT()), |
| StructType().add("key", LongType()).add("mat", MatrixUDT()), |
| ]: |
| cdf = self.connect.createDataFrame(data=[], schema=schema) |
| sdf = self.spark.createDataFrame(data=[], schema=schema) |
| |
| self.assertEqual(cdf.schema, sdf.schema) |
| |
| |
| if __name__ == "__main__": |
| from pyspark.sql.tests.connect.test_connect_creation import * # noqa: F401 |
| |
| try: |
| import xmlrunner |
| |
| testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) |
| except ImportError: |
| testRunner = None |
| |
| unittest.main(testRunner=testRunner, verbosity=2) |
