Refactor {year,month,day,hour} transform (#1563)

Similar to #1562
apache · Jan 23, 2025 · a0cc563 · a0cc563
1 parent 60d6e53
commit a0cc563
Showing 1 changed file with 21 additions and 0 deletions.
diff --git 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: