Skip to content

Commit c5f007b

Browse files
zhengruifengdongjoon-hyun
authored andcommitted
[SPARK-53829][PYTHON] Support datetime.time in column operators
### What changes were proposed in this pull request? Support `datetime.time` in column operators ### Why are the changes needed? to be consistent with other datetype types ### Does this PR introduce _any_ user-facing change? yes, this will work ```py df.select("i").where(sf.col("t") < datetime.time(3, 0, 0)) ``` ### How was this patch tested? new tests ### Was this patch authored or co-authored using generative AI tooling? no Closes #52540 from zhengruifeng/py_col_time. Authored-by: Ruifeng Zheng <[email protected]> Signed-off-by: Dongjoon Hyun <[email protected]>
1 parent 8bf6640 commit c5f007b

File tree

3 files changed

+35
-4
lines changed

3 files changed

+35
-4
lines changed

python/pyspark/sql/connect/_typing.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@
3939

4040
DecimalLiteral = decimal.Decimal
4141

42-
DateTimeLiteral = Union[datetime.datetime, datetime.date]
42+
DateTimeLiteral = Union[datetime.date, datetime.time, datetime.datetime]
4343

4444
DataTypeOrString = Union[DataType, str]
4545

python/pyspark/sql/connect/column.py

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -84,10 +84,11 @@ def _bin_op(
8484
float,
8585
int,
8686
str,
87-
datetime.datetime,
8887
datetime.date,
89-
decimal.Decimal,
88+
datetime.time,
89+
datetime.datetime,
9090
datetime.timedelta,
91+
decimal.Decimal,
9192
),
9293
):
9394
other_expr = LiteralExpression._from_value(other)
@@ -384,7 +385,17 @@ def substr(
384385
def __eq__(self, other: Any) -> ParentColumn: # type: ignore[override]
385386
other = enum_to_value(other)
386387
if other is None or isinstance(
387-
other, (bool, float, int, str, datetime.datetime, datetime.date, decimal.Decimal)
388+
other,
389+
(
390+
bool,
391+
float,
392+
int,
393+
str,
394+
datetime.date,
395+
datetime.time,
396+
datetime.datetime,
397+
decimal.Decimal,
398+
),
388399
):
389400
other_expr = LiteralExpression._from_value(other)
390401
else:

python/pyspark/sql/tests/test_column.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,26 @@ def test_column_operators(self):
113113
ValueError, "Cannot apply 'in' operator against a column", lambda: 1 in cs
114114
)
115115

116+
def test_column_date_time_op(self):
117+
query = """
118+
SELECT * FROM VALUES
119+
(TIME('00:00:00'), 1),
120+
(TIME('01:02:03'), 2),
121+
(TIME('11:12:13'), 3)
122+
AS tab(t, i)
123+
"""
124+
125+
df = self.spark.sql(query)
126+
127+
res1 = df.select("i").where(sf.col("t") < datetime.time(3, 0, 0))
128+
self.assertEqual([r.i for r in res1.collect()], [1, 2])
129+
130+
res2 = df.select("i").where(sf.col("t") > datetime.time(1, 0, 0))
131+
self.assertEqual([r.i for r in res2.collect()], [2, 3])
132+
133+
res3 = df.select("i").where(sf.col("t") == datetime.time(0, 0, 0))
134+
self.assertEqual([r.i for r in res3.collect()], [1])
135+
116136
def test_column_accessor(self):
117137
from pyspark.sql.functions import col
118138

0 commit comments

Comments
 (0)