Skip to content

Commit

Permalink
fix(pyspark): gate other usage of DayTimeIntervalType for PySpark 3.2
Browse files Browse the repository at this point in the history
  • Loading branch information
gforsyth committed Sep 7, 2023
1 parent ea3a090 commit ab01de0
Showing 1 changed file with 14 additions and 7 deletions.
21 changes: 14 additions & 7 deletions ibis/backends/pyspark/datatypes.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
from __future__ import annotations

import pyspark
import pyspark.sql.types as pt
from packaging.version import parse as vparse

import ibis.common.exceptions as com
import ibis.expr.datatypes as dt
Expand All @@ -9,6 +11,10 @@

_sql_type_names = dict(sql_type_names, date="date")

# DayTimeIntervalType introduced in Spark 3.2 (at least) but didn't show up in
# PySpark until version 3.3
PYSPARK_33 = vparse(pyspark.__version__) >= vparse("3.3")


def type_to_sql_string(tval):
if tval.is_decimal():
Expand Down Expand Up @@ -38,12 +44,13 @@ def type_to_sql_string(tval):
_to_pyspark_dtypes = {v: k for k, v in _from_pyspark_dtypes.items()}
_to_pyspark_dtypes[dt.JSON] = pt.StringType

_pyspark_interval_units = {
pt.DayTimeIntervalType.SECOND: "s",
pt.DayTimeIntervalType.MINUTE: "m",
pt.DayTimeIntervalType.HOUR: "h",
pt.DayTimeIntervalType.DAY: "D",
}
if PYSPARK_33:
_pyspark_interval_units = {
pt.DayTimeIntervalType.SECOND: "s",
pt.DayTimeIntervalType.MINUTE: "m",
pt.DayTimeIntervalType.HOUR: "h",
pt.DayTimeIntervalType.DAY: "D",
}


class PySparkType(TypeMapper):
Expand All @@ -64,7 +71,7 @@ def to_ibis(cls, typ, nullable=True):
fields = {f.name: cls.to_ibis(f.dataType) for f in typ.fields}

return dt.Struct(fields, nullable=nullable)
elif isinstance(typ, pt.DayTimeIntervalType):
elif PYSPARK_33 and isinstance(typ, pt.DayTimeIntervalType):
if (
typ.startField == typ.endField
and typ.startField in _pyspark_interval_units
Expand Down

0 comments on commit ab01de0

Please sign in to comment.