add timestamp_tz_ns support to the pyspark module

Author: evertlammerts

duckdb/experimental/spark/sql/type_utils.py

@@ -27,6 +27,7 @@
     TimeNTZType,
     TimestampMillisecondNTZType,
     TimestampNanosecondNTZType,
+    TimestampNanosecondType,
     TimestampNTZType,
     TimestampSecondNTZType,
     TimestampType,
@@ -62,9 +63,10 @@
     "time with time zone": TimeType,
     "timestamp": TimestampNTZType,
     "timestamp with time zone": TimestampType,
-    "timestamp_ms": TimestampNanosecondNTZType,
-    "timestamp_ns": TimestampMillisecondNTZType,
+    "timestamp_ms": TimestampMillisecondNTZType,
+    "timestamp_ns": TimestampNanosecondNTZType,
     "timestamp_s": TimestampSecondNTZType,
+    "timestamptz_ns": TimestampNanosecondType,
     "interval": DayTimeIntervalType,
     "list": ArrayType,
     "struct": StructType,

duckdb/experimental/spark/sql/types.py

@@ -49,6 +49,7 @@
     "TimestampMillisecondNTZType",
     "TimestampNTZType",
     "TimestampNanosecondNTZType",
+    "TimestampNanosecondType",
     "TimestampSecondNTZType",
     "TimestampType",
     "UUIDType",
@@ -239,6 +240,26 @@ def fromInternal(self, ts: int) -> datetime.datetime:  # noqa: D102
             return datetime.datetime.fromtimestamp(ts // 1000000).replace(microsecond=ts % 1000000)
 
 
+class TimestampNanosecondType(AtomicType, metaclass=DataTypeSingleton):
+    """Timestamp (datetime.datetime) data type with timezone information with nanosecond precision."""
+
+    def __init__(self) -> None:  # noqa: D107
+        super().__init__(DuckDBPyType("TIMESTAMPTZ_NS"))
+
+    def needConversion(self) -> bool:  # noqa: D102
+        return True
+
+    @classmethod
+    def typeName(cls) -> str:  # noqa: D102
+        return "timestamptz_ns"
+
+    def toInternal(self, dt: datetime.datetime) -> int:  # noqa: D102
+        raise ContributionsAcceptedError
+
+    def fromInternal(self, ts: int) -> datetime.datetime:  # noqa: D102
+        raise ContributionsAcceptedError
+
+
 class TimestampNTZType(AtomicType, metaclass=DataTypeSingleton):
     """Timestamp (datetime.datetime) data type without timezone information with microsecond precision."""
 

tests/fast/spark/test_spark_types.py

@@ -32,6 +32,7 @@
     TimeNTZType,
     TimestampMillisecondNTZType,
     TimestampNanosecondNTZType,
+    TimestampNanosecondType,
     TimestampNTZType,
     TimestampSecondNTZType,
     TimestampType,
@@ -86,10 +87,11 @@ def test_all_types_schema(self, spark):
                 StructField("time", TimeNTZType(), True),
                 StructField("timestamp", TimestampNTZType(), True),
                 StructField("timestamp_s", TimestampSecondNTZType(), True),
-                StructField("timestamp_ms", TimestampNanosecondNTZType(), True),
-                StructField("timestamp_ns", TimestampMillisecondNTZType(), True),
+                StructField("timestamp_ms", TimestampMillisecondNTZType(), True),
+                StructField("timestamp_ns", TimestampNanosecondNTZType(), True),
                 StructField("time_tz", TimeType(), True),
                 StructField("timestamp_tz", TimestampType(), True),
+                StructField("timestamp_tz_ns", TimestampNanosecondType(), True),
                 StructField("float", FloatType(), True),
                 StructField("double", DoubleType(), True),
                 StructField("dec_4_1", DecimalType(4, 1), True),