[iceberg] branch master updated: Python: Add support `date`, `time` and `datetime` in to_bytes (#8214)

[fo...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

fokko pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iceberg.git


The following commit(s) were added to refs/heads/master by this push:
     new b34f4f2356 Python: Add support `date`, `time` and `datetime` in to_bytes (#8214)
b34f4f2356 is described below

commit b34f4f2356b9ee3cc5b5b5ce0eff554dc5f862d0
Author: Rishabh Srivastava <ri...@gmail.com>
AuthorDate: Fri Aug 4 17:12:19 2023 +0530

    Python: Add support `date`, `time` and `datetime` in to_bytes (#8214)
    
    * added datetime support for to_bytes
    
    * fmt
    
    * resuing datetime_to_micros
    
    * removed unused imports
    
    * added new test cases
    
    * fixed format
    
    * removed type checking
    
    * added new test function
    
    * removed old test cases
    
    * datetypes added in function argument
    
    ---------
    
    Co-authored-by: Rishabh Srivastava <ri...@cmd.com.au>
---
 python/pyiceberg/conversions.py  | 30 ++++++++++++++++++++++++++----
 python/tests/test_conversions.py | 23 ++++++++++++++++++++++-
 2 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/python/pyiceberg/conversions.py b/python/pyiceberg/conversions.py
index 26ec71475d..ebf6cfa666 100644
--- a/python/pyiceberg/conversions.py
+++ b/python/pyiceberg/conversions.py
@@ -28,6 +28,7 @@ Note:
     implementations that share the same conversion logic, registrations can be stacked.
 """
 import uuid
+from datetime import date, datetime, time
 from decimal import Decimal
 from functools import singledispatch
 from struct import Struct
@@ -56,6 +57,7 @@ from pyiceberg.types import (
     TimeType,
     UUIDType,
 )
+from pyiceberg.utils.datetime import date_to_days, datetime_to_micros, time_object_to_micros
 from pyiceberg.utils.decimal import decimal_to_bytes, unscaled_to_decimal
 
 _BOOL_STRUCT = Struct("<?")
@@ -152,7 +154,9 @@ def _(_: DecimalType, value_str: str) -> Decimal:
 
 
 @singledispatch
-def to_bytes(primitive_type: PrimitiveType, _: Union[bool, bytes, Decimal, float, int, str, uuid.UUID]) -> bytes:
+def to_bytes(
+    primitive_type: PrimitiveType, _: Union[bool, bytes, Decimal, date, datetime, float, int, str, time, uuid.UUID]
+) -> bytes:
     """A generic function which converts a built-in python value to bytes.
 
     This conversion follows the serialization scheme for storing single values as individual binary values defined in the Iceberg specification that
@@ -172,16 +176,34 @@ def _(_: BooleanType, value: bool) -> bytes:
 
 
 @to_bytes.register(IntegerType)
-@to_bytes.register(DateType)
 def _(_: PrimitiveType, value: int) -> bytes:
     return _INT_STRUCT.pack(value)
 
 
 @to_bytes.register(LongType)
-@to_bytes.register(TimeType)
+def _(_: PrimitiveType, value: int) -> bytes:
+    return _LONG_STRUCT.pack(value)
+
+
 @to_bytes.register(TimestampType)
 @to_bytes.register(TimestamptzType)
-def _(_: PrimitiveType, value: int) -> bytes:
+def _(_: TimestampType, value: Union[datetime, int]) -> bytes:
+    if isinstance(value, datetime):
+        value = datetime_to_micros(value)
+    return _LONG_STRUCT.pack(value)
+
+
+@to_bytes.register(DateType)
+def _(_: DateType, value: Union[date, int]) -> bytes:
+    if isinstance(value, date):
+        value = date_to_days(value)
+    return _INT_STRUCT.pack(value)
+
+
+@to_bytes.register(TimeType)
+def _(_: TimeType, value: Union[time, int]) -> bytes:
+    if isinstance(value, time):
+        value = time_object_to_micros(value)
     return _LONG_STRUCT.pack(value)
 
 
diff --git a/python/tests/test_conversions.py b/python/tests/test_conversions.py
index 94f553741f..429de6e011 100644
--- a/python/tests/test_conversions.py
+++ b/python/tests/test_conversions.py
@@ -73,8 +73,14 @@ Notes:
 """
 import struct
 import uuid
+from datetime import (
+    date,
+    datetime,
+    time,
+    timezone,
+)
 from decimal import Decimal
-from typing import Any
+from typing import Any, Union
 
 import pytest
 
@@ -507,3 +513,18 @@ def test_raise_on_incorrect_precision_or_scale(primitive_type: DecimalType, valu
         conversions.to_bytes(primitive_type, value)
 
     assert expected_error_message in str(exc_info.value)
+
+
+@pytest.mark.parametrize(
+    "primitive_type, value, expected_bytes",
+    [
+        (TimestampType(), datetime(2023, 3, 1, 19, 25, 0), b"\x00\xbb\r\xab\xdb\xf5\x05\x00"),
+        (TimestamptzType(), datetime(2023, 3, 1, 19, 25, 0, tzinfo=timezone.utc), b"\x00\xbb\r\xab\xdb\xf5\x05\x00"),
+        (DateType(), date(2023, 3, 1), b"\xd9K\x00\x00"),
+        (TimeType(), time(12, 30, 45, 500000), b"`\xc8\xeb|\n\x00\x00\x00"),
+    ],
+)
+def test_datetime_obj_to_bytes(primitive_type: PrimitiveType, value: Union[datetime, date, time], expected_bytes: bytes) -> None:
+    bytes_from_value = conversions.to_bytes(primitive_type, value)
+
+    assert bytes_from_value == expected_bytes