From a0cc56351f724aa4c96d5548c79884f26ea16ce8 Mon Sep 17 00:00:00 2001
From: Fokko Driesprong <fokko@apache.org>
Date: Thu, 23 Jan 2025 05:04:48 +0100
Subject: [PATCH] Refactor `{year,month,day,hour}` transform (#1563)

Similar to https://github.com/apache/iceberg-python/pull/1562
---
 pyiceberg/transforms.py | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/pyiceberg/transforms.py b/pyiceberg/transforms.py
index 40a78b2811..b8f0b975e6 100644
--- a/pyiceberg/transforms.py
+++ b/pyiceberg/transforms.py
@@ -444,11 +444,17 @@ def transform(self, source: IcebergType) -> Callable[[Optional[S]], Optional[int
         if isinstance(source, DateType):
 
             def year_func(v: Any) -> int:
+                if isinstance(v, py_datetime.date):
+                    v = datetime.date_to_days(v)
+
                 return datetime.days_to_years(v)
 
         elif isinstance(source, (TimestampType, TimestamptzType)):
 
             def year_func(v: Any) -> int:
+                if isinstance(v, py_datetime.datetime):
+                    v = datetime.datetime_to_micros(v)
+
                 return datetime.micros_to_years(v)
 
         else:
@@ -501,11 +507,17 @@ def transform(self, source: IcebergType) -> Callable[[Optional[S]], Optional[int
         if isinstance(source, DateType):
 
             def month_func(v: Any) -> int:
+                if isinstance(v, py_datetime.date):
+                    v = datetime.date_to_days(v)
+
                 return datetime.days_to_months(v)
 
         elif isinstance(source, (TimestampType, TimestamptzType)):
 
             def month_func(v: Any) -> int:
+                if isinstance(v, py_datetime.datetime):
+                    v = datetime.datetime_to_micros(v)
+
                 return datetime.micros_to_months(v)
 
         else:
@@ -564,11 +576,17 @@ def transform(self, source: IcebergType) -> Callable[[Optional[S]], Optional[int
         if isinstance(source, DateType):
 
             def day_func(v: Any) -> int:
+                if isinstance(v, py_datetime.date):
+                    v = datetime.date_to_days(v)
+
                 return v
 
         elif isinstance(source, (TimestampType, TimestamptzType)):
 
             def day_func(v: Any) -> int:
+                if isinstance(v, py_datetime.datetime):
+                    v = datetime.datetime_to_micros(v)
+
                 return datetime.micros_to_days(v)
 
         else:
@@ -629,6 +647,9 @@ def transform(self, source: IcebergType) -> Callable[[Optional[S]], Optional[int
         if isinstance(source, (TimestampType, TimestamptzType)):
 
             def hour_func(v: Any) -> int:
+                if isinstance(v, py_datetime.datetime):
+                    v = datetime.datetime_to_micros(v)
+
                 return datetime.micros_to_hours(v)
 
         else: