add nanoseconds support to transforms

Author: sungwy

pyiceberg/transforms.py

@@ -73,7 +73,9 @@
     IntegerType,
     LongType,
     StringType,
+    TimestampNanoType,
     TimestampType,
+    TimestamptzNanoType,
     TimestamptzType,
     TimeType,
     UUIDType,
@@ -457,13 +459,20 @@ def year_func(v: Any) -> int:
 
                 return datetime.micros_to_years(v)
 
+        elif isinstance(source, (TimestampNanoType, TimestamptzNanoType)):
+
+            def year_func(v: Any) -> int:
+                # python datetime has no nanoseconds support.
+                # nanosecond datetimes will be expressed as int as a workaround
+                return datetime.nanos_to_years(v)
+
         else:
             raise ValueError(f"Cannot apply year transform for type: {source}")
 
         return lambda v: year_func(v) if v is not None else None
 
     def can_transform(self, source: IcebergType) -> bool:
-        return isinstance(source, (DateType, TimestampType, TimestamptzType))
+        return isinstance(source, (DateType, TimestampType, TimestamptzType, TimestampNanoType, TimestamptzNanoType))
 
     @property
     def granularity(self) -> TimeResolution:
@@ -481,15 +490,19 @@ def pyarrow_transform(self, source: IcebergType) -> "Callable[[pa.Array], pa.Arr
         import pyarrow.compute as pc
 
         if isinstance(source, DateType):
-            epoch = datetime.EPOCH_DATE
+            epoch = pa.scalar(datetime.EPOCH_DATE)
         elif isinstance(source, TimestampType):
-            epoch = datetime.EPOCH_TIMESTAMP
+            epoch = pa.scalar(datetime.EPOCH_TIMESTAMP)
         elif isinstance(source, TimestamptzType):
-            epoch = datetime.EPOCH_TIMESTAMPTZ
+            epoch = pa.scalar(datetime.EPOCH_TIMESTAMPTZ)
+        elif isinstance(source, TimestampNanoType):
+            epoch = pa.scalar(datetime.EPOCH_TIMESTAMP).cast(pa.timestamp("ns"))
+        elif isinstance(source, TimestamptzNanoType):
+            epoch = pa.scalar(datetime.EPOCH_TIMESTAMPTZ).cast(pa.timestamp("ns"))
         else:
             raise ValueError(f"Cannot apply year transform for type: {source}")
 
-        return lambda v: pc.years_between(pa.scalar(epoch), v) if v is not None else None
+        return lambda v: pc.years_between(epoch, v) if v is not None else None
 
 
 class MonthTransform(TimeTransform[S]):
@@ -520,13 +533,20 @@ def month_func(v: Any) -> int:
 
                 return datetime.micros_to_months(v)
 
+        elif isinstance(source, (TimestampNanoType, TimestamptzNanoType)):
+
+            def month_func(v: Any) -> int:
+                # python datetime has no nanoseconds support.
+                # nanosecond datetimes will be expressed as int as a workaround
+                return datetime.nanos_to_months(v)
+
         else:
             raise ValueError(f"Cannot apply month transform for type: {source}")
 
         return lambda v: month_func(v) if v is not None else None
 
     def can_transform(self, source: IcebergType) -> bool:
-        return isinstance(source, (DateType, TimestampType, TimestamptzType))
+        return isinstance(source, (DateType, TimestampType, TimestamptzType, TimestampNanoType, TimestamptzNanoType))
 
     @property
     def granularity(self) -> TimeResolution:
@@ -544,17 +564,21 @@ def pyarrow_transform(self, source: IcebergType) -> "Callable[[pa.Array], pa.Arr
         import pyarrow.compute as pc
 
         if isinstance(source, DateType):
-            epoch = datetime.EPOCH_DATE
+            epoch = pa.scalar(datetime.EPOCH_DATE)
         elif isinstance(source, TimestampType):
-            epoch = datetime.EPOCH_TIMESTAMP
+            epoch = pa.scalar(datetime.EPOCH_TIMESTAMP)
         elif isinstance(source, TimestamptzType):
-            epoch = datetime.EPOCH_TIMESTAMPTZ
+            epoch = pa.scalar(datetime.EPOCH_TIMESTAMPTZ)
+        elif isinstance(source, TimestampNanoType):
+            epoch = pa.scalar(datetime.EPOCH_TIMESTAMP).cast(pa.timestamp("ns"))
+        elif isinstance(source, TimestamptzNanoType):
+            epoch = pa.scalar(datetime.EPOCH_TIMESTAMPTZ).cast(pa.timestamp("ns"))
         else:
             raise ValueError(f"Cannot apply month transform for type: {source}")
 
         def month_func(v: pa.Array) -> pa.Array:
             return pc.add(
-                pc.multiply(pc.years_between(pa.scalar(epoch), v), pa.scalar(12)),
+                pc.multiply(pc.years_between(epoch, v), pa.scalar(12)),
                 pc.add(pc.month(v), pa.scalar(-1)),
             )
 
@@ -589,13 +613,20 @@ def day_func(v: Any) -> int:
 
                 return datetime.micros_to_days(v)
 
+        elif isinstance(source, (TimestampNanoType, TimestamptzNanoType)):
+
+            def day_func(v: Any) -> int:
+                # python datetime has no nanoseconds support.
+                # nanosecond datetimes will be expressed as int as a workaround
+                return datetime.nanos_to_days(v)
+
         else:
             raise ValueError(f"Cannot apply day transform for type: {source}")
 
         return lambda v: day_func(v) if v is not None else None
 
     def can_transform(self, source: IcebergType) -> bool:
-        return isinstance(source, (DateType, TimestampType, TimestamptzType))
+        return isinstance(source, (DateType, TimestampType, TimestamptzType, TimestampNanoType, TimestamptzNanoType))
 
     def result_type(self, source: IcebergType) -> IcebergType:
         """Return the result type of a day transform.
@@ -621,15 +652,19 @@ def pyarrow_transform(self, source: IcebergType) -> "Callable[[pa.Array], pa.Arr
         import pyarrow.compute as pc
 
         if isinstance(source, DateType):
-            epoch = datetime.EPOCH_DATE
+            epoch = pa.scalar(datetime.EPOCH_DATE)
         elif isinstance(source, TimestampType):
-            epoch = datetime.EPOCH_TIMESTAMP
+            epoch = pa.scalar(datetime.EPOCH_TIMESTAMP)
         elif isinstance(source, TimestamptzType):
-            epoch = datetime.EPOCH_TIMESTAMPTZ
+            epoch = pa.scalar(datetime.EPOCH_TIMESTAMPTZ)
+        elif isinstance(source, TimestampNanoType):
+            epoch = pa.scalar(datetime.EPOCH_TIMESTAMP).cast(pa.timestamp("ns"))
+        elif isinstance(source, TimestamptzNanoType):
+            epoch = pa.scalar(datetime.EPOCH_TIMESTAMPTZ).cast(pa.timestamp("ns"))
         else:
             raise ValueError(f"Cannot apply day transform for type: {source}")
 
-        return lambda v: pc.days_between(pa.scalar(epoch), v) if v is not None else None
+        return lambda v: pc.days_between(epoch, v) if v is not None else None
 
 
 class HourTransform(TimeTransform[S]):
@@ -652,13 +687,20 @@ def hour_func(v: Any) -> int:
 
                 return datetime.micros_to_hours(v)
 
+        elif isinstance(source, (TimestampNanoType, TimestamptzNanoType)):
+
+            def day_func(v: Any) -> int:
+                # python datetime has no nanoseconds support.
+                # nanosecond datetimes will be expressed as int as a workaround
+                return datetime.nanos_to_hours(v)
+
         else:
             raise ValueError(f"Cannot apply hour transform for type: {source}")
 
         return lambda v: hour_func(v) if v is not None else None
 
     def can_transform(self, source: IcebergType) -> bool:
-        return isinstance(source, (TimestampType, TimestamptzType))
+        return isinstance(source, (TimestampType, TimestamptzType, TimestampNanoType, TimestamptzNanoType))
 
     @property
     def granularity(self) -> TimeResolution:
@@ -676,13 +718,17 @@ def pyarrow_transform(self, source: IcebergType) -> "Callable[[pa.Array], pa.Arr
         import pyarrow.compute as pc
 
         if isinstance(source, TimestampType):
-            epoch = datetime.EPOCH_TIMESTAMP
+            epoch = pa.scalar(datetime.EPOCH_TIMESTAMP)
         elif isinstance(source, TimestamptzType):
-            epoch = datetime.EPOCH_TIMESTAMPTZ
+            epoch = pa.scalar(datetime.EPOCH_TIMESTAMPTZ)
+        elif isinstance(source, TimestampNanoType):
+            epoch = pa.scalar(datetime.EPOCH_TIMESTAMP).cast(pa.timestamp("ns"))
+        elif isinstance(source, TimestamptzNanoType):
+            epoch = pa.scalar(datetime.EPOCH_TIMESTAMPTZ).cast(pa.timestamp("ns"))
         else:
             raise ValueError(f"Cannot apply hour transform for type: {source}")
 
-        return lambda v: pc.hours_between(pa.scalar(epoch), v) if v is not None else None
+        return lambda v: pc.hours_between(epoch, v) if v is not None else None
 
 
 def _base64encode(buffer: bytes) -> str:

pyiceberg/utils/datetime.py

@@ -217,3 +217,38 @@ def days_to_years(days: int) -> int:
 
 def micros_to_years(micros: int) -> int:
     return micros_to_timestamp(micros).year - EPOCH_TIMESTAMP.year
+
+
+def nanos_to_timestamp(nanos: int) -> datetime:
+    """Convert nanoseconds from epoch to a microsecond timestamp."""
+    dt = timedelta(microseconds=nanos // 1000)
+    return EPOCH_TIMESTAMP + dt
+
+
+def nanos_to_years(nanos: int) -> int:
+    return nanos_to_timestamp(nanos).year - EPOCH_TIMESTAMP.year
+
+
+def nanos_to_months(nanos: int) -> int:
+    dt = nanos_to_timestamp(nanos)
+    return (dt.year - EPOCH_TIMESTAMP.year) * 12 + (dt.month - EPOCH_TIMESTAMP.month)
+
+
+def nanos_to_days(nanos: int) -> int:
+    """Convert a timestamp in nanoseconds to a date in days."""
+    return timedelta(microseconds=nanos // 1000).days
+
+
+def nanos_to_time(nanos: int) -> time:
+    """Convert a timestamp in nanoseconds to a microsecond precision time."""
+    micros = nanos // 1000
+    micros, microseconds = divmod(micros, 1000000)
+    micros, seconds = divmod(micros, 60)
+    micros, minutes = divmod(micros, 60)
+    hours = micros
+    return time(hour=hours, minute=minutes, second=seconds, microsecond=microseconds)
+
+
+def nanos_to_hours(nanos: int) -> int:
+    """Convert a timestamp in nanoseconds to hours from 1970-01-01T00:00."""
+    return nanos // 3_600_000_000_0000

tests/test_transforms.py

@@ -1559,7 +1559,7 @@ def test_ymd_pyarrow_transforms(
         ]
     else:
         with pytest.raises(ValueError):
-            transform.pyarrow_transform(DateType())(arrow_table_date_timestamps[source_col])
+            transform.pyarrow_transform(source_type)(arrow_table_date_timestamps[source_col])
 
 
 @pytest.mark.parametrize(

tests/utils/test_datetime.py

@@ -80,7 +80,7 @@ def test_millis_to_datetime() -> None:
     assert millis_to_datetime(1690971805918) == datetime(2023, 8, 2, 10, 23, 25, 918000)
 
 
-@pytest.mark.parametrize("time_str, nanos", [("00:00:00Z", 0), ("20:21:44.375612-0500", 73304375612000)])
+@pytest.mark.parametrize("time_str, nanos", [("00:00:00Z", 0), ("20:21:44.375612-05:00", 73304375612000)])
 def test_time_str_to_nanos(time_str: str, nanos: int) -> None:
     assert nanos == time_str_to_nanos(time_str)