| # -*- 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. |
| # |
| |
| from enum import Enum |
| from itertools import chain |
| import datetime |
| import unittest |
| |
| from pyspark.sql import Column, Row |
| from pyspark.sql import functions as sf |
| from pyspark.sql.types import StructType, StructField, IntegerType, LongType |
| from pyspark.errors import AnalysisException, PySparkTypeError, PySparkValueError |
| from pyspark.testing.sqlutils import ReusedSQLTestCase, have_pandas, pandas_requirement_message |
| |
| |
| class ColumnTestsMixin: |
| def test_column_name_encoding(self): |
| """Ensure that created columns has `str` type consistently.""" |
| columns = self.spark.createDataFrame([("Alice", 1)], ["name", "age"]).columns |
| self.assertEqual(columns, ["name", "age"]) |
| self.assertTrue(isinstance(columns[0], str)) |
| self.assertTrue(isinstance(columns[1], str)) |
| |
| 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_validate_column_types(self): |
| from pyspark.sql.functions import udf, to_json |
| from pyspark.sql.classic.column import _to_java_column |
| |
| self.assertTrue("Column" in _to_java_column("a").getClass().toString()) |
| self.assertTrue("Column" in _to_java_column("a").getClass().toString()) |
| self.assertTrue("Column" in _to_java_column(self.spark.range(1).id).getClass().toString()) |
| |
| with self.assertRaises(PySparkTypeError) as pe: |
| _to_java_column(1) |
| |
| self.check_error( |
| exception=pe.exception, |
| errorClass="NOT_COLUMN_OR_STR", |
| messageParameters={"arg_name": "col", "arg_type": "int"}, |
| ) |
| |
| class A: |
| pass |
| |
| self.assertRaises(TypeError, lambda: _to_java_column(A())) |
| self.assertRaises(TypeError, lambda: _to_java_column([])) |
| |
| with self.assertRaises(PySparkTypeError) as pe: |
| udf(lambda x: x)(None) |
| |
| self.check_error( |
| exception=pe.exception, |
| errorClass="NOT_COLUMN_OR_STR", |
| messageParameters={"arg_name": "col", "arg_type": "NoneType"}, |
| ) |
| self.assertRaises(TypeError, lambda: to_json(1)) |
| |
| def test_column_operators(self): |
| ci = self.df.key |
| cs = self.df.value |
| 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.contains("a"), |
| cs.like("a"), |
| cs.rlike("a"), |
| cs.ilike("A"), |
| cs.asc(), |
| cs.desc(), |
| cs.startswith("a"), |
| cs.endswith("a"), |
| ci.eqNullSafe(cs), |
| sf.col("b") & sf.lit(True), |
| sf.col("b") & True, |
| sf.lit(True) & sf.col("b"), |
| True & sf.col("b"), |
| sf.col("b") | sf.lit(True), |
| sf.col("b") | True, |
| sf.lit(True) | sf.col("b"), |
| True | sf.col("b"), |
| ) |
| self.assertTrue(all(isinstance(c, Column) for c in css)) |
| self.assertTrue(isinstance(ci.cast(LongType()), Column)) |
| self.assertRaisesRegex( |
| ValueError, "Cannot apply 'in' operator against a column", lambda: 1 in cs |
| ) |
| |
| def test_column_accessor(self): |
| from pyspark.sql.functions import col |
| |
| self.assertIsInstance(col("foo")[1:3], Column) |
| self.assertIsInstance(col("foo")[0], Column) |
| self.assertIsInstance(col("foo")["bar"], Column) |
| self.assertRaises(ValueError, lambda: col("foo")[0:10:2]) |
| |
| 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_access_column(self): |
| df = self.df |
| self.assertTrue(isinstance(df.key, Column)) |
| self.assertTrue(isinstance(df["key"], Column)) |
| self.assertTrue(isinstance(df[0], Column)) |
| self.assertRaises(IndexError, lambda: df[2]) |
| self.assertRaises(TypeError, lambda: df[{}]) |
| self.assertRaises(AnalysisException, lambda: df.select(df["bad_key"]).schema) |
| |
| def test_column_name_with_non_ascii(self): |
| columnName = "数量" |
| self.assertTrue(isinstance(columnName, str)) |
| schema = StructType([StructField(columnName, LongType(), True)]) |
| df = self.spark.createDataFrame([(1,)], schema) |
| self.assertEqual(schema, df.schema) |
| self.assertEqual("DataFrame[数量: bigint]", str(df)) |
| self.assertEqual([("数量", "bigint")], df.dtypes) |
| self.assertEqual(1, df.select("数量").first()[0]) |
| self.assertEqual(1, df.select(df["数量"]).first()[0]) |
| self.assertTrue(columnName in repr(df[columnName])) |
| |
| def test_field_accessor(self): |
| df = self.spark.createDataFrame([Row(l=[1], r=Row(a=1, b="b"), d={"k": "v"})]) |
| self.assertEqual(1, df.select(df.l[0]).first()[0]) |
| self.assertEqual(1, df.select(df.r["a"]).first()[0]) |
| self.assertEqual(1, df.select(df["r.a"]).first()[0]) |
| self.assertEqual("b", df.select(df.r["b"]).first()[0]) |
| self.assertEqual("b", df.select(df["r.b"]).first()[0]) |
| self.assertEqual("v", df.select(df.d["k"]).first()[0]) |
| |
| def test_bitwise_operations(self): |
| from pyspark.sql import functions |
| |
| row = Row(a=170, b=75) |
| df = self.spark.createDataFrame([row]) |
| result = df.select(df.a.bitwiseAND(df.b)).collect()[0].asDict() |
| self.assertEqual(170 & 75, result["(a & b)"]) |
| result = df.select(df.a.bitwiseOR(df.b)).collect()[0].asDict() |
| self.assertEqual(170 | 75, result["(a | b)"]) |
| result = df.select(df.a.bitwiseXOR(df.b)).collect()[0].asDict() |
| self.assertEqual(170 ^ 75, result["(a ^ b)"]) |
| result = df.select(functions.bitwiseNOT(df.b)).collect()[0].asDict() |
| self.assertEqual(~75, result["~b"]) |
| result = df.select(functions.bitwise_not(df.b)).collect()[0].asDict() |
| self.assertEqual(~75, result["~b"]) |
| |
| def test_with_field(self): |
| from pyspark.sql.functions import lit, col |
| |
| df = self.spark.createDataFrame([Row(a=Row(b=1, c=2))]) |
| self.assertIsInstance(df["a"].withField("b", lit(3)), Column) |
| self.assertIsInstance(df["a"].withField("d", lit(3)), Column) |
| result = df.withColumn("a", df["a"].withField("d", lit(3))).collect()[0].asDict() |
| self.assertEqual(3, result["a"]["d"]) |
| result = df.withColumn("a", df["a"].withField("b", lit(3))).collect()[0].asDict() |
| self.assertEqual(3, result["a"]["b"]) |
| |
| with self.assertRaises(PySparkTypeError) as pe: |
| df["a"].withField("b", 3) |
| |
| self.check_error( |
| exception=pe.exception, |
| errorClass="NOT_COLUMN", |
| messageParameters={"arg_name": "col", "arg_type": "int"}, |
| ) |
| |
| with self.assertRaises(PySparkTypeError) as pe: |
| df["a"].withField(col("b"), lit(3)) |
| |
| self.check_error( |
| exception=pe.exception, |
| errorClass="NOT_STR", |
| messageParameters={"arg_name": "fieldName", "arg_type": "Column"}, |
| ) |
| |
| def test_drop_fields(self): |
| df = self.spark.createDataFrame([Row(a=Row(b=1, c=2, d=Row(e=3, f=4)))]) |
| self.assertIsInstance(df["a"].dropFields("b"), Column) |
| self.assertIsInstance(df["a"].dropFields("b", "c"), Column) |
| self.assertIsInstance(df["a"].dropFields("d.e"), Column) |
| |
| result = ( |
| df.select( |
| df["a"].dropFields("b").alias("a1"), |
| df["a"].dropFields("d.e").alias("a2"), |
| ) |
| .first() |
| .asDict(True) |
| ) |
| |
| self.assertTrue("b" not in result["a1"] and "c" in result["a1"] and "d" in result["a1"]) |
| |
| self.assertTrue("e" not in result["a2"]["d"] and "f" in result["a2"]["d"]) |
| |
| def test_getitem_column(self): |
| mapping = {"A": "20", "B": "28", "C": "34"} |
| mapping_expr = sf.create_map([sf.lit(x) for x in chain(*mapping.items())]) |
| df = self.spark.createDataFrame( |
| data=[["A", "10"], ["B", "14"], ["C", "17"]], |
| schema=["key", "value"], |
| ).withColumn("square_value", mapping_expr[sf.col("key")]) |
| self.assertEqual(df.count(), 3) |
| |
| def test_alias_negative(self): |
| with self.assertRaises(PySparkValueError) as pe: |
| self.spark.range(1).id.alias("a", "b", metadata={}) |
| |
| self.check_error( |
| exception=pe.exception, |
| errorClass="ONLY_ALLOWED_FOR_SINGLE_COLUMN", |
| messageParameters={"arg_name": "metadata"}, |
| ) |
| |
| def test_cast_str_representation(self): |
| self.assertEqual(str(sf.col("a").cast("int")), "Column<'CAST(a AS INT)'>") |
| self.assertEqual(str(sf.col("a").cast("INT")), "Column<'CAST(a AS INT)'>") |
| self.assertEqual(str(sf.col("a").cast(IntegerType())), "Column<'CAST(a AS INT)'>") |
| self.assertEqual(str(sf.col("a").cast(LongType())), "Column<'CAST(a AS BIGINT)'>") |
| |
| self.assertEqual(str(sf.col("a").try_cast("int")), "Column<'TRY_CAST(a AS INT)'>") |
| self.assertEqual(str(sf.col("a").try_cast("INT")), "Column<'TRY_CAST(a AS INT)'>") |
| self.assertEqual(str(sf.col("a").try_cast(IntegerType())), "Column<'TRY_CAST(a AS INT)'>") |
| self.assertEqual(str(sf.col("a").try_cast(LongType())), "Column<'TRY_CAST(a AS BIGINT)'>") |
| |
| def test_cast_negative(self): |
| with self.assertRaises(PySparkTypeError) as pe: |
| self.spark.range(1).id.cast(123) |
| |
| self.check_error( |
| exception=pe.exception, |
| errorClass="NOT_DATATYPE_OR_STR", |
| messageParameters={"arg_name": "dataType", "arg_type": "int"}, |
| ) |
| |
| def test_over_negative(self): |
| with self.assertRaises(PySparkTypeError) as pe: |
| self.spark.range(1).id.over(123) |
| |
| self.check_error( |
| exception=pe.exception, |
| errorClass="NOT_WINDOWSPEC", |
| messageParameters={"arg_name": "window", "arg_type": "int"}, |
| ) |
| |
| def test_eqnullsafe_classmethod_usage(self): |
| df = self.spark.range(1) |
| self.assertEqual(df.select(Column.eqNullSafe(df.id, df.id)).first()[0], True) |
| |
| def test_isinstance_dataframe(self): |
| self.assertIsInstance(self.spark.range(1).id, Column) |
| |
| def test_expr_str_representation(self): |
| expression = sf.expr("foo") |
| when_cond = sf.when(expression, sf.lit(None)) |
| self.assertEqual(str(when_cond), "Column<'CASE WHEN foo THEN NULL END'>") |
| |
| def test_col_field_ops_representation(self): |
| # SPARK-49894: Test string representation of columns |
| c = sf.col("c") |
| |
| # getField |
| self.assertEqual(str(c.x), "Column<'c['x']'>") |
| self.assertEqual(str(c.x.y), "Column<'c['x']['y']'>") |
| self.assertEqual(str(c.x.y.z), "Column<'c['x']['y']['z']'>") |
| |
| self.assertEqual(str(c["x"]), "Column<'c['x']'>") |
| self.assertEqual(str(c["x"]["y"]), "Column<'c['x']['y']'>") |
| self.assertEqual(str(c["x"]["y"]["z"]), "Column<'c['x']['y']['z']'>") |
| |
| self.assertEqual(str(c.getField("x")), "Column<'c['x']'>") |
| self.assertEqual( |
| str(c.getField("x").getField("y")), |
| "Column<'c['x']['y']'>", |
| ) |
| self.assertEqual( |
| str(c.getField("x").getField("y").getField("z")), |
| "Column<'c['x']['y']['z']'>", |
| ) |
| |
| self.assertEqual(str(c.getItem("x")), "Column<'c['x']'>") |
| self.assertEqual( |
| str(c.getItem("x").getItem("y")), |
| "Column<'c['x']['y']'>", |
| ) |
| self.assertEqual( |
| str(c.getItem("x").getItem("y").getItem("z")), |
| "Column<'c['x']['y']['z']'>", |
| ) |
| |
| self.assertEqual( |
| str(c.x["y"].getItem("z")), |
| "Column<'c['x']['y']['z']'>", |
| ) |
| self.assertEqual( |
| str(c["x"].getField("y").getItem("z")), |
| "Column<'c['x']['y']['z']'>", |
| ) |
| self.assertEqual( |
| str(c.getField("x").getItem("y").z), |
| "Column<'c['x']['y']['z']'>", |
| ) |
| self.assertEqual( |
| str(c["x"].y.getField("z")), |
| "Column<'c['x']['y']['z']'>", |
| ) |
| |
| # WithField |
| self.assertEqual( |
| str(c.withField("x", sf.col("y"))), |
| "Column<'update_field(c, x, y)'>", |
| ) |
| self.assertEqual( |
| str(c.withField("x", sf.col("y")).withField("x", sf.col("z"))), |
| "Column<'update_field(update_field(c, x, y), x, z)'>", |
| ) |
| |
| # DropFields |
| self.assertEqual(str(c.dropFields("x")), "Column<'drop_field(c, x)'>") |
| self.assertEqual( |
| str(c.dropFields("x", "y")), |
| "Column<'drop_field(drop_field(c, x), y)'>", |
| ) |
| self.assertEqual( |
| str(c.dropFields("x", "y", "z")), |
| "Column<'drop_field(drop_field(drop_field(c, x), y), z)'>", |
| ) |
| |
| def test_lit_time_representation(self): |
| dt = datetime.date(2021, 3, 4) |
| self.assertEqual(str(sf.lit(dt)), "Column<'2021-03-04'>") |
| |
| ts = datetime.datetime(2021, 3, 4, 12, 34, 56, 1234) |
| self.assertEqual(str(sf.lit(ts)), "Column<'2021-03-04 12:34:56.001234'>") |
| |
| ts = datetime.time(12, 34, 56, 1234) |
| self.assertEqual(str(sf.lit(ts)), "Column<'12:34:56.001234'>") |
| |
| @unittest.skipIf(not have_pandas, pandas_requirement_message) |
| def test_lit_delta_representation(self): |
| for delta in [ |
| datetime.timedelta(days=1), |
| datetime.timedelta(hours=2), |
| datetime.timedelta(minutes=3), |
| datetime.timedelta(seconds=4), |
| datetime.timedelta(microseconds=5), |
| datetime.timedelta(days=2, hours=21, microseconds=908), |
| datetime.timedelta(days=1, minutes=-3, microseconds=-1001), |
| datetime.timedelta(days=1, hours=2, minutes=3, seconds=4, microseconds=5), |
| ]: |
| import pandas as pd |
| |
| # Column<'PT69H0.000908S'> or Column<'P2DT21H0M0.000908S'> |
| s = str(sf.lit(delta)) |
| |
| # Parse the ISO string representation and compare |
| self.assertTrue(pd.Timedelta(s[8:-2]).to_pytimedelta() == delta) |
| |
| def test_enum_literals(self): |
| class IntEnum(Enum): |
| X = 1 |
| Y = 2 |
| Z = 3 |
| |
| class BoolEnum(Enum): |
| T = True |
| |
| class StrEnum(Enum): |
| X = "x" |
| |
| id = sf.col("id") |
| s = sf.col("s") |
| b = sf.col("b") |
| |
| cols, expected = list( |
| zip( |
| (id + IntEnum.X, 2), |
| (id - IntEnum.X, 0), |
| (id * IntEnum.X, 1), |
| (id / IntEnum.X, 1.0), |
| (id % IntEnum.X, 0), |
| (IntEnum.X + id, 2), |
| (IntEnum.X - id, 0), |
| (IntEnum.X * id, 1), |
| (IntEnum.X / id, 1.0), |
| (IntEnum.X % id, 0), |
| (id**IntEnum.X, 1.0), |
| (IntEnum.X**id, 1, 0), |
| (id == IntEnum.X, True), |
| (IntEnum.X == id, True), |
| (id < IntEnum.X, False), |
| (id <= IntEnum.X, True), |
| (id >= IntEnum.X, True), |
| (id > IntEnum.X, False), |
| (id.eqNullSafe(IntEnum.X), True), |
| (b & BoolEnum.T, True), |
| (b | BoolEnum.T, True), |
| (BoolEnum.T & b, True), |
| (BoolEnum.T | b, True), |
| (id.bitwiseOR(IntEnum.X), 1), |
| (id.bitwiseAND(IntEnum.X), 1), |
| (id.bitwiseXOR(IntEnum.X), 0), |
| (id.contains(IntEnum.X), True), |
| (s.startswith(StrEnum.X), False), |
| (s.endswith(StrEnum.X), False), |
| (s.like(StrEnum.X), False), |
| (s.rlike(StrEnum.X), False), |
| (s.ilike(StrEnum.X), False), |
| (s.substr(IntEnum.X, IntEnum.Y), "1"), |
| (sf.when(b, IntEnum.X).when(~b, IntEnum.Y).otherwise(IntEnum.Z), 1), |
| ) |
| ) |
| result = ( |
| self.spark.range(1, 2) |
| .select(id, id.astype("string").alias("s"), id.astype("boolean").alias("b")) |
| .select(*cols) |
| .first() |
| ) |
| |
| for r, c, e in zip(result, cols, expected): |
| self.assertEqual(r, e, str(c)) |
| |
| |
| class ColumnTests(ColumnTestsMixin, ReusedSQLTestCase): |
| pass |
| |
| |
| if __name__ == "__main__": |
| import unittest |
| from pyspark.sql.tests.test_column 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) |
