/
githubmirror
/
spark
Обзор
Документация
Войти
/
githubmirror
/
spark
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
master
python/pyspark/sql/tests/test_dataframe.py
1 234 строки
49 KB
Nikolaus Schuetz
[MINOR][PYTHON] Fix wrong expected_type in DataFrame.sort ascending type error
02 авг 2026, 20:27
02 авг 2026, 20:27
1222a59
Код
Авторство
О чём код?
# # 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 glob import os import pydoc import shutil import tempfile import warnings import unittest import io from contextlib import redirect_stdout from pyspark.sql import Row, functions, DataFrame from pyspark.sql.functions import ( col, lit, count, struct, date_format, to_date, array, explode, ) from pyspark.sql.types import ( StringType, IntegerType, LongType, StructType, StructField, ) from pyspark.storagelevel import StorageLevel from pyspark.errors import ( AnalysisException, IllegalArgumentException, PySparkTypeError, PySparkValueError, ) from pyspark.testing import assertDataFrameEqual from pyspark.testing.sqlutils import ( ReusedSQLTestCase, SPARK_HOME, ) from pyspark.testing.utils import ( have_pyarrow, have_pandas, pandas_requirement_message, pyarrow_requirement_message, ) class DataFrameTestsMixin: 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_table(self): with self.assertRaises(PySparkTypeError) as pe: self.spark.table(None) self.check_error( exception=pe.exception, errorClass="NOT_EXPECTED_TYPE", messageParameters={ "expected_type": "str", "arg_name": "tableName", "arg_type": "NoneType", }, ) def test_dataframe_star(self): df1 = self.spark.createDataFrame([{"a": 1}]) df2 = self.spark.createDataFrame([{"a": 1, "b": "v"}]) df3 = df2.withColumnsRenamed({"a": "x", "b": "y"}) df = df1.join(df2) self.assertEqual(df.columns, ["a", "a", "b"]) self.assertEqual(df.select(df1["*"]).columns, ["a"]) self.assertEqual(df.select(df2["*"]).columns, ["a", "b"]) df = df1.join(df2).withColumn("c", lit(0)) self.assertEqual(df.columns, ["a", "a", "b", "c"]) self.assertEqual(df.select(df1["*"]).columns, ["a"]) self.assertEqual(df.select(df2["*"]).columns, ["a", "b"]) df = df1.join(df2, "a") self.assertEqual(df.columns, ["a", "b"]) self.assertEqual(df.select(df1["*"]).columns, ["a"]) self.assertEqual(df.select(df2["*"]).columns, ["a", "b"]) df = df1.join(df2, "a").withColumn("c", lit(0)) self.assertEqual(df.columns, ["a", "b", "c"]) self.assertEqual(df.select(df1["*"]).columns, ["a"]) self.assertEqual(df.select(df2["*"]).columns, ["a", "b"]) df = df2.join(df3) self.assertEqual(df.columns, ["a", "b", "x", "y"]) self.assertEqual(df.select(df2["*"]).columns, ["a", "b"]) self.assertEqual(df.select(df3["*"]).columns, ["x", "y"]) df = df2.join(df3).withColumn("c", lit(0)) self.assertEqual(df.columns, ["a", "b", "x", "y", "c"]) self.assertEqual(df.select(df2["*"]).columns, ["a", "b"]) self.assertEqual(df.select(df3["*"]).columns, ["x", "y"]) def test_count_star(self): df1 = self.spark.createDataFrame([{"a": 1}]) df2 = self.spark.createDataFrame([{"a": 1, "b": "v"}]) df3 = df2.select(struct("a", "b").alias("s")) self.assertEqual(df1.select(count(df1["*"])).columns, ["count(1)"]) self.assertEqual(df1.select(count(col("*"))).columns, ["count(1)"]) self.assertEqual(df2.select(count(df2["*"])).columns, ["count(1)"]) self.assertEqual(df2.select(count(col("*"))).columns, ["count(1)"]) self.assertEqual(df3.select(count(df3["*"])).columns, ["count(1)"]) self.assertEqual(df3.select(count(col("*"))).columns, ["count(1)"]) def test_lateral_column_alias(self): df1 = self.spark.range(10).select( (col("id") + lit(1)).alias("x"), (col("x") + lit(1)).alias("y") ) df2 = self.spark.range(10).select(col("id").alias("x")) df3 = df1.join(df2, df1.x == df2.x).select(df1.y) self.assertEqual(df3.columns, ["y"]) self.assertEqual(df3.count(), 9) 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_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 self.check_help_command(df) def check_help_command(self, df): pydoc.render_doc(df) pydoc.render_doc(df.foo) pydoc.render_doc(df.take(1)) def test_drop(self): df = self.spark.createDataFrame([("A", 50, "Y"), ("B", 60, "Y")], ["name", "age", "active"]) self.assertEqual(df.drop("active").columns, ["name", "age"]) self.assertEqual(df.drop("active", "nonexistent_column").columns, ["name", "age"]) self.assertEqual(df.drop("name", "age", "active").columns, []) self.assertEqual(df.drop(col("name")).columns, ["age", "active"]) self.assertEqual(df.drop(col("name"), col("age")).columns, ["active"]) self.assertEqual(df.drop(col("name"), col("age"), col("random")).columns, ["active"]) def test_drop_col_from_different_dataframe(self): df1 = self.spark.range(10) df2 = df1.withColumn("v0", lit(0)) # drop df2["id"] from df2 self.assertEqual(df2.drop(df2["id"]).columns, ["v0"]) # drop df1["id"] from df2, which is semantically equal to df2["id"] # note that df1.drop(df2["id"]) works in Classic, but not in Connect self.assertEqual(df2.drop(df1["id"]).columns, ["v0"]) df3 = df2.select("*", lit(1).alias("v1")) # drop df3["id"] from df3 self.assertEqual(df3.drop(df3["id"]).columns, ["v0", "v1"]) # drop df2["id"] from df3, which is semantically equal to df3["id"] self.assertEqual(df3.drop(df2["id"]).columns, ["v0", "v1"]) # drop df1["id"] from df3, which is semantically equal to df3["id"] self.assertEqual(df3.drop(df1["id"]).columns, ["v0", "v1"]) # drop df3["v0"] from df3 self.assertEqual(df3.drop(df3["v0"]).columns, ["id", "v1"]) # drop df2["v0"] from df3, which is semantically equal to df3["v0"] self.assertEqual(df3.drop(df2["v0"]).columns, ["id", "v1"]) def test_drop_join(self): left_df = self.spark.createDataFrame( [(1, "a"), (2, "b"), (3, "c")], ["join_key", "value1"], ) right_df = self.spark.createDataFrame( [(1, "aa"), (2, "bb"), (4, "dd")], ["join_key", "value2"], ) joined_df = left_df.join( right_df, on=left_df["join_key"] == right_df["join_key"], how="left", ) dropped_1 = joined_df.drop(left_df["join_key"]) self.assertEqual(dropped_1.columns, ["value1", "join_key", "value2"]) self.assertEqual( dropped_1.sort("value1").collect(), [ Row(value1="a", join_key=1, value2="aa"), Row(value1="b", join_key=2, value2="bb"), Row(value1="c", join_key=None, value2=None), ], ) dropped_2 = joined_df.drop(right_df["join_key"]) self.assertEqual(dropped_2.columns, ["join_key", "value1", "value2"]) self.assertEqual( dropped_2.sort("value1").collect(), [ Row(join_key=1, value1="a", value2="aa"), Row(join_key=2, value1="b", value2="bb"), Row(join_key=3, value1="c", value2=None), ], ) def test_with_columns_renamed(self): df = self.spark.createDataFrame([("Alice", 50), ("Alice", 60)], ["name", "age"]) # rename both columns renamed_df1 = df.withColumnsRenamed({"name": "naam", "age": "leeftijd"}) self.assertEqual(renamed_df1.columns, ["naam", "leeftijd"]) # rename one column with one missing name renamed_df2 = df.withColumnsRenamed({"name": "naam", "address": "adres"}) self.assertEqual(renamed_df2.columns, ["naam", "age"]) # negative test for incorrect type with self.assertRaises(PySparkTypeError) as pe: df.withColumnsRenamed(("name", "x")) self.check_error( exception=pe.exception, errorClass="NOT_EXPECTED_TYPE", messageParameters={"expected_type": "dict", "arg_name": "colsMap", "arg_type": "tuple"}, ) def test_sort_ascending_invalid_type(self): df = self.spark.createDataFrame([("Alice", 10)], ["name", "age"]) with self.assertRaises(PySparkTypeError) as pe: df.sort("age", ascending="asc") self.check_error( exception=pe.exception, errorClass="NOT_EXPECTED_TYPE", messageParameters={ "expected_type": "bool, int or list", "arg_name": "ascending", "arg_type": "str", }, ) def test_with_columns_renamed_with_duplicated_names(self): df1 = self.spark.createDataFrame([(1, "v1")], ["id", "value"]) df2 = self.spark.createDataFrame([(1, "x", "v2")], ["id", "a", "value"]) join = df2.join(df1, on=["id"], how="left") self.assertEqual( join.withColumnRenamed("id", "value").columns, join.withColumnsRenamed({"id": "value"}).columns, ) self.assertEqual( join.withColumnRenamed("a", "b").columns, join.withColumnsRenamed({"a": "b"}).columns, ) self.assertEqual( join.withColumnRenamed("value", "new_value").columns, join.withColumnsRenamed({"value": "new_value"}).columns, ) self.assertEqual( join.withColumnRenamed("x", "y").columns, join.withColumnsRenamed({"x": "y"}).columns, ) def test_ordering_of_with_columns_renamed(self): df = self.spark.range(10) df1 = df.withColumnsRenamed({"id": "a", "a": "b"}) self.assertEqual(df1.columns, ["b"]) df2 = df.withColumnsRenamed({"a": "b", "id": "a"}) self.assertEqual(df2.columns, ["a"]) def test_drop_duplicates(self): # SPARK-36034 test that drop duplicates throws a type error when in correct type provided df = self.spark.createDataFrame([("Alice", 50), ("Alice", 60)], ["name", "age"]) # shouldn't drop a non-null row self.assertEqual(df.dropDuplicates().count(), 2) self.assertEqual(df.dropDuplicates(["name"]).count(), 1) self.assertEqual(df.dropDuplicates(["name", "age"]).count(), 2) with self.assertRaises(PySparkTypeError) as pe: df.dropDuplicates("name") self.check_error( exception=pe.exception, errorClass="NOT_EXPECTED_TYPE", messageParameters={ "expected_type": "list or tuple", "arg_name": "subset", "arg_type": "str", }, ) # Should raise proper error when taking non-string values with self.assertRaises(PySparkTypeError) as pe: df.dropDuplicates([None]).show() self.check_error( exception=pe.exception, errorClass="NOT_EXPECTED_TYPE", messageParameters={ "expected_type": "str", "arg_name": "subset", "arg_type": "NoneType", }, ) with self.assertRaises(PySparkTypeError) as pe: df.dropDuplicates([1]).show() self.check_error( exception=pe.exception, errorClass="NOT_EXPECTED_TYPE", messageParameters={"expected_type": "str", "arg_name": "subset", "arg_type": "int"}, ) def test_drop_duplicates_with_ambiguous_reference(self): df1 = self.spark.createDataFrame([(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) df2 = self.spark.createDataFrame([Row(height=80, name="Tom"), Row(height=85, name="Bob")]) df3 = df1.join(df2, df1.name == df2.name, "inner") self.assertEqual(df3.drop("name", "age").columns, ["height"]) self.assertEqual(df3.drop("name", df3.age, "unknown").columns, ["height"]) self.assertEqual(df3.drop("name", "age", df3.height).columns, []) def test_drop_empty_column(self): df = self.spark.createDataFrame([(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) self.assertEqual(df.drop().columns, ["age", "name"]) self.assertEqual(df.drop(*[]).columns, ["age", "name"]) def test_drop_column_name_with_dot(self): df = ( self.spark.range(1, 3) .withColumn("first.name", lit("Peter")) .withColumn("city.name", lit("raleigh")) .withColumn("state", lit("nc")) ) self.assertEqual(df.drop("first.name").columns, ["id", "city.name", "state"]) self.assertEqual(df.drop("city.name").columns, ["id", "first.name", "state"]) self.assertEqual(df.drop("first.name", "city.name").columns, ["id", "state"]) self.assertEqual( df.drop("first.name", "city.name", "unknown.unknown").columns, ["id", "state"] ) self.assertEqual( df.drop("unknown.unknown").columns, ["id", "first.name", "city.name", "state"] ) 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_with_columns(self): # With single column keys = self.df.withColumns({"key": self.df.key}).select("key").collect() self.assertEqual([r.key for r in keys], list(range(100))) # With key and value columns kvs = ( self.df.withColumns({"key": self.df.key, "value": self.df.value}) .select("key", "value") .collect() ) self.assertEqual([(r.key, r.value) for r in kvs], [(i, str(i)) for i in range(100)]) # Columns rename kvs = ( self.df.withColumns({"key_alias": self.df.key, "value_alias": self.df.value}) .select("key_alias", "value_alias") .collect() ) self.assertEqual( [(r.key_alias, r.value_alias) for r in kvs], [(i, str(i)) for i in range(100)] ) # Type check self.assertRaises(TypeError, self.df.withColumns, ["key"]) self.assertRaises(Exception, self.df.withColumns) def test_generic_hints(self): df1 = self.spark.range(10e10).toDF("id") df2 = self.spark.range(10e10).toDF("id") self.assertIsInstance(df1.hint("broadcast"), type(df1)) # Dummy rules self.assertIsInstance(df1.hint("broadcast", "foo", "bar"), type(df1)) with io.StringIO() as buf, redirect_stdout(buf): df1.join(df2.hint("broadcast"), "id").explain(True) self.assertEqual(1, buf.getvalue().count("BroadcastHashJoin")) def test_coalesce_hints_with_string_parameter(self): with self.sql_conf({"spark.sql.adaptive.coalescePartitions.enabled": False}): df = self.spark.createDataFrame( zip(["A", "B"] * 2**9, range(2**10)), StructType([StructField("a", StringType()), StructField("n", IntegerType())]), ) with io.StringIO() as buf, redirect_stdout(buf): # COALESCE coalesce = df.hint("coalesce", 2) coalesce.explain(True) output = buf.getvalue() self.assertGreaterEqual(output.count("Coalesce 2"), 1) buf.truncate(0) buf.seek(0) # REPARTITION_BY_RANGE range_partitioned = df.hint("REPARTITION_BY_RANGE", 2, "a") range_partitioned.explain(True) output = buf.getvalue() self.assertGreaterEqual(output.count("REPARTITION_BY_NUM"), 1) buf.truncate(0) buf.seek(0) # REBALANCE rebalanced1 = df.hint("REBALANCE", "a") # just check this doesn't error rebalanced1.explain(True) rebalanced2 = df.hint("REBALANCE", 2) rebalanced2.explain(True) rebalanced3 = df.hint("REBALANCE", 2, "a") rebalanced3.explain(True) rebalanced4 = df.hint("REBALANCE", functions.col("a")) rebalanced4.explain(True) output = buf.getvalue() self.assertGreaterEqual(output.count("REBALANCE_PARTITIONS_BY_NONE"), 1) self.assertGreaterEqual(output.count("REBALANCE_PARTITIONS_BY_COL"), 3) # add tests for SPARK-23647 (test more types for hint) def test_extended_hint_types(self): df = self.spark.range(10e10).toDF("id") such_a_nice_list = ["itworks1", "itworks2", "itworks3"] int_list = [1, 2, 3] hinted_df = df.hint("my awesome hint", 1.2345, "what", such_a_nice_list, int_list) self.assertIsInstance(df.hint("broadcast", []), type(df)) self.assertIsInstance(df.hint("broadcast", ["foo", "bar"]), type(df)) with io.StringIO() as buf, redirect_stdout(buf): hinted_df.explain(True) explain_output = buf.getvalue() self.assertGreaterEqual(explain_output.count("1.2345"), 1) self.assertGreaterEqual(explain_output.count("what"), 1) self.assertGreaterEqual(explain_output.count("itworks"), 1) def test_sample(self): with self.assertRaises(PySparkTypeError) as pe: self.spark.range(1).sample() self.check_error( exception=pe.exception, errorClass="NOT_EXPECTED_TYPE", messageParameters={ "expected_type": "bool, float or int", "arg_name": "withReplacement (optional), fraction (required) and seed (optional)", "arg_type": "NoneType, NoneType, NoneType", }, ) 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).count() ) def test_sample_with_random_seed(self): df = self.spark.range(10000).sample(0.1) cnts = [df.count() for i in range(10)] self.assertEqual(1, len(set(cnts))) def test_toDF_with_string(self): df = self.spark.createDataFrame([("John", 30), ("Alice", 25), ("Bob", 28)]) data = [("John", 30), ("Alice", 25), ("Bob", 28)] result = df.toDF("key", "value") self.assertEqual(result.schema.simpleString(), "struct<key:string,value:bigint>") self.assertEqual(result.collect(), data) with self.assertRaises(PySparkTypeError) as pe: df.toDF("key", None) self.check_error( exception=pe.exception, errorClass="NOT_EXPECTED_TYPE", messageParameters={ "expected_type": "list[str]", "arg_name": "cols", "arg_type": "NoneType", }, ) 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.assertRaisesRegex( Exception, "FIELD_STRUCT_LENGTH_MISMATCH", lambda: rdd.coalesce(1).toDF("key: int").collect(), ) # field types mismatch will cause exception at runtime. self.assertRaisesRegex( Exception, "FIELD_DATA_TYPE_UNACCEPTABLE", lambda: rdd.coalesce(1).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_create_df_with_collation(self): schema = StructType([StructField("name", StringType("UNICODE_CI"), True)]) df = self.spark.createDataFrame([("Alice",), ("alice",)], schema) self.assertEqual(df.select("name").distinct().count(), 1) def test_print_schema(self): df = self.spark.createDataFrame([(1, (2, 2))], ["a", "b"]) with io.StringIO() as buf, redirect_stdout(buf): df.printSchema(1) self.assertEqual(1, buf.getvalue().count("long")) self.assertEqual(0, buf.getvalue().count("_1")) self.assertEqual(0, buf.getvalue().count("_2")) buf.truncate(0) buf.seek(0) df.printSchema(2) self.assertEqual(3, buf.getvalue().count("long")) self.assertEqual(1, buf.getvalue().count("_1")) self.assertEqual(1, buf.getvalue().count("_2")) 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(AnalysisException, 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")) with self.sql_conf({"spark.sql.crossJoin.enabled": False}): # 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_dataframe(self): df = self.spark.createDataFrame([(2, 2), (3, 3)]) try: self.assertEqual(df.storageLevel, StorageLevel.NONE) df.cache() self.assertEqual(df.storageLevel, StorageLevel.MEMORY_AND_DISK_DESER) df.unpersist() self.assertEqual(df.storageLevel, StorageLevel.NONE) df.persist() self.assertEqual(df.storageLevel, StorageLevel.MEMORY_AND_DISK_DESER) df.unpersist(blocking=True) self.assertEqual(df.storageLevel, StorageLevel.NONE) df.persist(StorageLevel.DISK_ONLY) self.assertEqual(df.storageLevel, StorageLevel.DISK_ONLY) finally: df.unpersist() self.assertEqual(df.storageLevel, StorageLevel.NONE) def test_cache_table(self): spark = self.spark tables = ["tab1", "tab2", "tab3"] with self.temp_view(*tables): for i, tab in enumerate(tables): spark.createDataFrame([(2, i), (3, i)]).createOrReplaceTempView(tab) self.assertFalse(spark.catalog.isCached(tab)) spark.catalog.cacheTable("tab1") spark.catalog.cacheTable("tab3", StorageLevel.OFF_HEAP) self.assertTrue(spark.catalog.isCached("tab1")) self.assertFalse(spark.catalog.isCached("tab2")) self.assertTrue(spark.catalog.isCached("tab3")) spark.catalog.cacheTable("tab2") spark.catalog.uncacheTable("tab1") spark.catalog.uncacheTable("tab3") self.assertFalse(spark.catalog.isCached("tab1")) self.assertTrue(spark.catalog.isCached("tab2")) self.assertFalse(spark.catalog.isCached("tab3")) spark.catalog.clearCache() self.assertFalse(spark.catalog.isCached("tab1")) self.assertFalse(spark.catalog.isCached("tab2")) self.assertFalse(spark.catalog.isCached("tab3")) self.assertRaisesRegex( AnalysisException, "does_not_exist", lambda: spark.catalog.isCached("does_not_exist"), ) self.assertRaisesRegex( AnalysisException, "does_not_exist", lambda: spark.catalog.cacheTable("does_not_exist"), ) self.assertRaisesRegex( AnalysisException, "does_not_exist", lambda: spark.catalog.uncacheTable("does_not_exist"), ) 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.assertEqual(re.sub(pattern, "", expected1), df.__repr__()) with self.sql_conf({"spark.sql.repl.eagerEval.truncate": 3}): expected2 = """+---+-----+ ||key|value| |+---+-----+ || 1| 1| ||222| 222| |+---+-----+ |""" self.assertEqual(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.assertEqual(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.assertEqual(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.assertEqual(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.assertEqual(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.assertEqual(None, df._repr_html_()) self.assertEqual(expected, df.__repr__()) with self.sql_conf({"spark.sql.repl.eagerEval.truncate": 3}): self.assertEqual(None, df._repr_html_()) self.assertEqual(expected, df.__repr__()) with self.sql_conf({"spark.sql.repl.eagerEval.maxNumRows": 1}): self.assertEqual(None, df._repr_html_()) self.assertEqual(expected, df.__repr__()) def test_same_semantics_error(self): with self.assertRaises(PySparkTypeError) as pe: self.spark.range(10).sameSemantics(1) self.check_error( exception=pe.exception, errorClass="NOT_EXPECTED_TYPE", messageParameters={ "expected_type": "DataFrame", "arg_name": "other", "arg_type": "int", }, ) def test_input_files(self): tpath = tempfile.mkdtemp() shutil.rmtree(tpath) try: self.spark.range(1, 100, 1, 10).write.parquet(tpath) # read parquet file and get the input files list input_files_list = self.spark.read.parquet(tpath).inputFiles() # input files list should contain 10 entries self.assertEqual(len(input_files_list), 10) # all file paths in list must contain tpath for file_path in input_files_list: self.assertTrue(tpath in file_path) finally: shutil.rmtree(tpath) def test_df_show(self): # SPARK-35408: ensure better diagnostics if incorrect parameters are passed # to DataFrame.show df = self.spark.createDataFrame([("foo",)]) df.show(5) df.show(5, True) df.show(5, 1, True) df.show(n=5, truncate="1", vertical=False) df.show(n=5, truncate=1.5, vertical=False) with self.assertRaises(PySparkTypeError) as pe: df.show(True) self.check_error( exception=pe.exception, errorClass="NOT_EXPECTED_TYPE", messageParameters={"expected_type": "int", "arg_name": "n", "arg_type": "bool"}, ) with self.assertRaises(PySparkTypeError) as pe: df.show(vertical="foo") self.check_error( exception=pe.exception, errorClass="NOT_EXPECTED_TYPE", messageParameters={"expected_type": "bool", "arg_name": "vertical", "arg_type": "str"}, ) with self.assertRaises(PySparkTypeError) as pe: df.show(truncate="foo") self.check_error( exception=pe.exception, errorClass="NOT_EXPECTED_TYPE", messageParameters={"expected_type": "bool", "arg_name": "truncate", "arg_type": "str"}, ) def test_df_merge_into(self): filename_pattern = ( "sql/catalyst/target/scala-*/test-classes/org/apache/spark/sql/connector/catalog/" "InMemoryRowLevelOperationTableCatalog.class" ) if not bool(glob.glob(os.path.join(SPARK_HOME, filename_pattern))): raise unittest.SkipTest( "org.apache.spark.sql.connector.catalog.InMemoryRowLevelOperationTableCatalog' " "is not available. Will skip the related tests" ) try: # InMemoryRowLevelOperationTableCatalog is a test catalog that is included in the # catalyst-test package. If Spark complains that it can't find this class, make sure # the catalyst-test JAR does exist in "SPARK_HOME/assembly/target/scala-x.xx/jars" # directory. If not, build it with `build/sbt test:assembly` and copy it over. self.spark.conf.set( "spark.sql.catalog.testcat", "org.apache.spark.sql.connector.catalog.InMemoryRowLevelOperationTableCatalog", ) with self.table("testcat.ns1.target"): def reset_target_table(): self.spark.createDataFrame( [(1, "Alice"), (2, "Bob")], ["id", "name"] ).write.mode("overwrite").saveAsTable("testcat.ns1.target") source = self.spark.createDataFrame([(1, "Charlie"), (3, "David")], ["id", "name"]) # type: DataFrame from pyspark.sql.functions import col # Match -> update, NotMatch -> insert, NotMatchedBySource -> delete reset_target_table() # fmt: off source.mergeInto("testcat.ns1.target", source.id == col("target.id")) \ .whenMatched(source.id == 1).update({"name": source.name}) \ .whenNotMatched().insert({"id": source.id, "name": source.name}) \ .whenNotMatchedBySource().delete() \ .merge() # fmt: on self.assertEqual( self.spark.table("testcat.ns1.target").orderBy("id").collect(), [Row(id=1, name="Charlie"), Row(id=3, name="David")], ) # Match -> updateAll, NotMatch -> insertAll, NotMatchedBySource -> update reset_target_table() # fmt: off source.mergeInto("testcat.ns1.target", source.id == col("target.id")) \ .whenMatched(source.id == 1).updateAll() \ .whenNotMatched(source.id == 3).insertAll() \ .whenNotMatchedBySource(col("target.id") == lit(2)) \ .update({"name": lit("not_matched")}) \ .merge() # fmt: on self.assertEqual( self.spark.table("testcat.ns1.target").orderBy("id").collect(), [ Row(id=1, name="Charlie"), Row(id=2, name="not_matched"), Row(id=3, name="David"), ], ) # Match -> delete, NotMatchedBySource -> delete reset_target_table() # fmt: off self.spark.createDataFrame([(1, "AliceJr")], ["id", "name"]) \ .write.mode("append").saveAsTable("testcat.ns1.target") source.mergeInto("testcat.ns1.target", source.id == col("target.id")) \ .whenMatched(col("target.name") != lit("AliceJr")).delete() \ .whenNotMatchedBySource().delete() \ .merge() # fmt: on self.assertEqual( self.spark.table("testcat.ns1.target").orderBy("id").collect(), [Row(id=1, name="AliceJr")], ) finally: self.spark.conf.unset("spark.sql.catalog.testcat") @unittest.skipIf( not have_pandas or not have_pyarrow, pandas_requirement_message or pyarrow_requirement_message, ) def test_pandas_api(self): import pandas as pd from pandas.testing import assert_frame_equal sdf = self.spark.createDataFrame([("a", 1), ("b", 2), ("c", 3)], ["Col1", "Col2"]) psdf_from_sdf = sdf.pandas_api() psdf_from_sdf_with_index = sdf.pandas_api(index_col="Col1") pdf = pd.DataFrame({"Col1": ["a", "b", "c"], "Col2": [1, 2, 3]}) pdf_with_index = pdf.set_index("Col1") assert_frame_equal(pdf, psdf_from_sdf.to_pandas()) assert_frame_equal(pdf_with_index, psdf_from_sdf_with_index.to_pandas()) def test_to(self): schema = StructType( [StructField("i", StringType(), True), StructField("j", IntegerType(), True)] ) df = self.spark.createDataFrame([("a", 1)], schema) schema1 = StructType([StructField("j", StringType()), StructField("i", StringType())]) df1 = df.to(schema1) self.assertEqual(schema1, df1.schema) self.assertEqual(df.count(), df1.count()) schema2 = StructType([StructField("j", LongType())]) df2 = df.to(schema2) self.assertEqual(schema2, df2.schema) self.assertEqual(df.count(), df2.count()) schema3 = StructType([StructField("struct", schema1, False)]) df3 = df.select(struct("i", "j").alias("struct")).to(schema3) self.assertEqual(schema3, df3.schema) self.assertEqual(df.count(), df3.count()) # incompatible field nullability schema4 = StructType([StructField("j", LongType(), False)]) self.assertRaisesRegex( AnalysisException, "NULLABLE_COLUMN_OR_FIELD", lambda: df.to(schema4).count() ) # field cannot upcast schema5 = StructType([StructField("i", LongType())]) self.assertRaisesRegex( AnalysisException, "INVALID_COLUMN_OR_FIELD_DATA_TYPE", lambda: df.to(schema5).count() ) def test_colregex(self): with self.assertRaises(PySparkTypeError) as pe: self.spark.range(10).colRegex(10) self.check_error( exception=pe.exception, errorClass="NOT_EXPECTED_TYPE", messageParameters={"expected_type": "str", "arg_name": "colName", "arg_type": "int"}, ) def test_where(self): with self.assertRaises(PySparkTypeError) as pe: self.spark.range(10).where(10) self.check_error( exception=pe.exception, errorClass="NOT_EXPECTED_TYPE", messageParameters={ "expected_type": "Column or str", "arg_name": "condition", "arg_type": "int", }, ) def test_duplicate_field_names(self): data = [ Row(Row("a", 1), Row(2, 3, "b", 4, "c", "d")), Row(Row("w", 6), Row(7, 8, "x", 9, "y", "z")), ] schema = ( StructType() .add("struct", StructType().add("x", StringType()).add("x", IntegerType())) .add( "struct", StructType() .add("a", IntegerType()) .add("x", IntegerType()) .add("x", StringType()) .add("y", IntegerType()) .add("y", StringType()) .add("x", StringType()), ) ) df = self.spark.createDataFrame(data, schema=schema) self.assertEqual(df.schema, schema) self.assertEqual(df.collect(), data) def test_union_classmethod_usage(self): df = self.spark.range(1) self.assertEqual(DataFrame.union(df, df).collect(), [Row(id=0), Row(id=0)]) def test_isinstance_dataframe(self): self.assertIsInstance(self.spark.range(1), DataFrame) def test_local_checkpoint_dataframe(self): with io.StringIO() as buf, redirect_stdout(buf): self.spark.range(1).localCheckpoint().explain() self.assertIn("ExistingRDD", buf.getvalue()) def test_local_checkpoint_dataframe_with_storage_level(self): # We don't have a way to reach into the server and assert the storage level server side, but # this test should cover for unexpected errors in the API. df = self.spark.range(10).localCheckpoint(eager=True, storageLevel=StorageLevel.DISK_ONLY) df.collect() def test_socket_leak(self): with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always", ResourceWarning) df = self.spark.range(10) df.collect() df = self.spark.range(10) for _ in df.toLocalIterator(): pass self.assertEqual(w, []) def test_transpose(self): df = self.spark.createDataFrame([{"a": "x", "b": "y", "c": "z"}]) # default index column transposed_df = df.transpose() expected_schema = StructType( [StructField("key", StringType(), False), StructField("x", StringType(), True)] ) expected_data = [Row(key="b", x="y"), Row(key="c", x="z")] expected_df = self.spark.createDataFrame(expected_data, schema=expected_schema) assertDataFrameEqual(transposed_df, expected_df, checkRowOrder=True) # specified index column transposed_df = df.transpose("c") expected_schema = StructType( [StructField("key", StringType(), False), StructField("z", StringType(), True)] ) expected_data = [Row(key="a", z="x"), Row(key="b", z="y")] expected_df = self.spark.createDataFrame(expected_data, schema=expected_schema) assertDataFrameEqual(transposed_df, expected_df, checkRowOrder=True) # enforce transpose max values with self.sql_conf({"spark.sql.transposeMaxValues": 0}): with self.assertRaises(AnalysisException) as pe: df.transpose().collect() self.check_error( exception=pe.exception, errorClass="TRANSPOSE_EXCEED_ROW_LIMIT", messageParameters={"maxValues": "0", "config": "spark.sql.transposeMaxValues"}, ) # enforce ascending order based on index column values for transposed columns df = self.spark.createDataFrame([{"a": "z"}, {"a": "y"}, {"a": "x"}]) transposed_df = df.transpose() expected_schema = StructType( [ StructField("key", StringType(), False), StructField("x", StringType(), True), StructField("y", StringType(), True), StructField("z", StringType(), True), ] ) # z, y, x -> x, y, z expected_df = self.spark.createDataFrame([], schema=expected_schema) assertDataFrameEqual(transposed_df, expected_df, checkRowOrder=True) # enforce AtomicType Attribute for index column values df = self.spark.createDataFrame([{"a": ["x", "x"], "b": "y", "c": "z"}]) with self.assertRaises(AnalysisException) as pe: df.transpose().collect() self.check_error( exception=pe.exception, errorClass="TRANSPOSE_INVALID_INDEX_COLUMN", messageParameters={ "reason": "Index column must be of atomic type, " "but found: ArrayType(StringType,true)" }, ) # enforce least common type for non-index columns df = self.spark.createDataFrame([{"a": "x", "b": "y", "c": 1}]) with self.assertRaises(AnalysisException) as pe: df.transpose().collect() self.check_error( exception=pe.exception, errorClass="TRANSPOSE_NO_LEAST_COMMON_TYPE", messageParameters={"dt1": '"STRING"', "dt2": '"BIGINT"'}, ) def test_transpose_with_invalid_index_columns(self): # SPARK-50602: invalid index columns df = self.spark.createDataFrame([{"a": "x", "b": "y", "c": "z"}]) with self.assertRaises(AnalysisException) as pe: df.transpose(col("a") + 1).collect() self.check_error( exception=pe.exception, errorClass="TRANSPOSE_INVALID_INDEX_COLUMN", messageParameters={"reason": "Index column must be an atomic attribute"}, ) def test_metadata_column(self): with self.sql_conf( {"spark.sql.catalog.testcat": "org.apache.spark.sql.connector.catalog.InMemoryCatalog"} ): tbl = "testcat.t" with self.table(tbl): self.spark.sql(f""" CREATE TABLE {tbl} (index bigint, data string) PARTITIONED BY (bucket(4, index), index) """) self.spark.sql(f"""INSERT INTO {tbl} VALUES (1, 'a'), (2, 'b'), (3, 'c')""") df = self.spark.sql(f"""SELECT * FROM {tbl}""") assertDataFrameEqual( df.select(df.metadataColumn("index")), [Row(0), Row(0), Row(0)], ) def test_with_column_and_generator(self): # SPARK-51451: Generators should be available with withColumn df = self.spark.createDataFrame([("082017",)], ["dt"]).select( to_date(col("dt"), "MMyyyy").alias("dt") ) df_dt = df.withColumn("dt", date_format(col("dt"), "MM/dd/yyyy")) monthArray = [lit(x) for x in range(0, 12)] df_month_y = df_dt.withColumn("month_y", explode(array(monthArray))) assertDataFrameEqual( df_month_y, [Row(dt="08/01/2017", month_y=i) for i in range(12)], ) df_dt_month_y = df.withColumns( { "dt": date_format(col("dt"), "MM/dd/yyyy"), "month_y": explode(array(monthArray)), } ) assertDataFrameEqual( df_dt_month_y, [Row(dt="08/01/2017", month_y=i) for i in range(12)], ) def test_query_execution_unsupported_in_classic(self): with self.assertRaises(PySparkValueError) as pe: self.spark.range(1).executionInfo self.check_error( exception=pe.exception, errorClass="CLASSIC_OPERATION_NOT_SUPPORTED_ON_DF", messageParameters={"member": "queryExecution"}, ) def test_to_json(self): df = self.spark.range(10) with self.sql_conf({"spark.sql.pyspark.toJSON.returnDataFrame": False}): from pyspark import RDD rdd = df.toJSON() self.assertIsInstance(rdd, RDD) self.assertEqual(rdd.count(), 10) with self.sql_conf({"spark.sql.pyspark.toJSON.returnDataFrame": True}): df = df.toJSON() self.assertIsInstance(df, DataFrame) self.assertEqual(df.select("value").count(), 10) def test_zip_with_index(self): df = self.spark.createDataFrame([("a", 1), ("b", 2), ("c", 3)], ["letter", "number"]) # Default column name "index" result = df.zipWithIndex() self.assertEqual(result.columns, ["letter", "number", "index"]) rows = result.collect() self.assertEqual(len(rows), 3) indices = [row["index"] for row in rows] self.assertEqual(sorted(indices), [0, 1, 2]) # Custom column name result = df.zipWithIndex("row_id") self.assertEqual(result.columns, ["letter", "number", "row_id"]) rows = result.collect() indices = [row["row_id"] for row in rows] self.assertEqual(sorted(indices), [0, 1, 2]) # Duplicate column name causes AMBIGUOUS_REFERENCE on select result = df.zipWithIndex("letter") with self.assertRaises(AnalysisException) as ctx: result.select("letter").collect() self.assertEqual(ctx.exception.getCondition(), "AMBIGUOUS_REFERENCE") # Duplicate column name causes COLUMN_ALREADY_EXISTS on write with tempfile.TemporaryDirectory() as d: with self.assertRaises(AnalysisException) as ctx: result.write.parquet(d) self.assertEqual(ctx.exception.getCondition(), "COLUMN_ALREADY_EXISTS") class DataFrameTests(DataFrameTestsMixin, ReusedSQLTestCase): pass if __name__ == "__main__": from pyspark.testing import main main()