Skip to content

Commit

Permalink
fix(pyspark): gate datediff op to restore pyspark 3.2 support
Browse files Browse the repository at this point in the history
  • Loading branch information
gforsyth authored and jcrist committed Sep 7, 2023
1 parent 993d902 commit 4a8d611
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 7 deletions.
21 changes: 14 additions & 7 deletions ibis/backends/pyspark/compiler.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
import pyspark
import pyspark.sql.functions as F
import pyspark.sql.types as pt
from packaging.version import parse as vparse
from pyspark.sql import Window
from pyspark.sql.functions import PandasUDFType, pandas_udf

Expand Down Expand Up @@ -1557,14 +1558,20 @@ def compile_date_sub(t, op, **kwargs):
)


@compiles(ops.DateDiff)
def compile_date_diff(t, op, **kwargs):
left = t.translate(op.left, **kwargs)
right = t.translate(op.right, **kwargs)
if vparse(pyspark.__version__) >= vparse("3.3"):

return F.concat(F.lit("INTERVAL '"), F.datediff(left, right), F.lit("' DAY")).cast(
pt.DayTimeIntervalType(pt.DayTimeIntervalType.DAY, pt.DayTimeIntervalType.DAY)
)
@compiles(ops.DateDiff)
def compile_date_diff(t, op, **kwargs):
left = t.translate(op.left, **kwargs)
right = t.translate(op.right, **kwargs)

return F.concat(
F.lit("INTERVAL '"), F.datediff(left, right), F.lit("' DAY")
).cast(
pt.DayTimeIntervalType(
pt.DayTimeIntervalType.DAY, pt.DayTimeIntervalType.DAY
)
)


@compiles(ops.TimestampAdd)
Expand Down
5 changes: 5 additions & 0 deletions ibis/backends/tests/test_temporal.py
Original file line number Diff line number Diff line change
Expand Up @@ -949,6 +949,11 @@ def convert_to_offset(x):
),
id="date-subtract-date",
marks=[
pytest.mark.xfail_version(
pyspark=["pyspark<3.3"],
raises=AttributeError,
reason="DayTimeIntervalType added in pyspark 3.3",
),
pytest.mark.notimpl(["bigquery"], raises=com.OperationNotDefinedError),
pytest.mark.notimpl(
["druid"],
Expand Down

0 comments on commit 4a8d611

Please sign in to comment.