From d38e3d496f0082b15c8eb0ab45a9b26c8f51d25a Mon Sep 17 00:00:00 2001
From: Prashanth Rao <35005448+prrao87@users.noreply.github.com>
Date: Mon, 16 Oct 2023 22:42:24 -0400
Subject: [PATCH] Add pyarrow date and timestamp type conversion from pydantic
 (#576)

---
 python/lancedb/pydantic.py    |  5 +++++
 python/tests/test_pydantic.py | 25 +++++++++++++++++++++++--
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/python/lancedb/pydantic.py b/python/lancedb/pydantic.py
index abe97068..2616f52f 100644
--- a/python/lancedb/pydantic.py
+++ b/python/lancedb/pydantic.py
@@ -19,6 +19,7 @@ import inspect
 import sys
 import types
 from abc import ABC, abstractmethod
+from datetime import date, datetime
 from typing import Any, Callable, Dict, Generator, List, Type, Union, _GenericAlias
 
 import numpy as np
@@ -159,6 +160,10 @@ def _py_type_to_arrow_type(py_type: Type[Any]) -> pa.DataType:
         return pa.bool_()
     elif py_type == bytes:
         return pa.binary()
+    elif py_type == date:
+        return pa.date32()
+    elif py_type == datetime:
+        return pa.timestamp("us")
     raise TypeError(
         f"Converting Pydantic type to Arrow Type: unsupported type {py_type}"
     )
diff --git a/python/tests/test_pydantic.py b/python/tests/test_pydantic.py
index 4da0e5ca..fa7e4849 100644
--- a/python/tests/test_pydantic.py
+++ b/python/tests/test_pydantic.py
@@ -14,6 +14,7 @@
 
 import json
 import sys
+from datetime import date, datetime
 from typing import List, Optional
 
 import pyarrow as pa
@@ -40,10 +41,18 @@ def test_pydantic_to_arrow():
         li: List[int]
         opt: Optional[str] = None
         st: StructModel
+        dt: date
+        dtt: datetime
         # d: dict
 
     m = TestModel(
-        id=1, s="hello", vec=[1.0, 2.0, 3.0], li=[2, 3, 4], st=StructModel(a="a", b=1.0)
+        id=1,
+        s="hello",
+        vec=[1.0, 2.0, 3.0],
+        li=[2, 3, 4],
+        st=StructModel(a="a", b=1.0),
+        dt=date.today(),
+        dtt=datetime.now(),
     )
 
     schema = pydantic_to_schema(TestModel)
@@ -62,6 +71,8 @@ def test_pydantic_to_arrow():
                 ),
                 False,
             ),
+            pa.field("dt", pa.date32(), False),
+            pa.field("dtt", pa.timestamp("us"), False),
         ]
     )
     assert schema == expect_schema
@@ -79,10 +90,18 @@ def test_pydantic_to_arrow_py38():
         li: List[int]
         opt: Optional[str] = None
         st: StructModel
+        dt: date
+        dtt: datetime
         # d: dict
 
     m = TestModel(
-        id=1, s="hello", vec=[1.0, 2.0, 3.0], li=[2, 3, 4], st=StructModel(a="a", b=1.0)
+        id=1,
+        s="hello",
+        vec=[1.0, 2.0, 3.0],
+        li=[2, 3, 4],
+        st=StructModel(a="a", b=1.0),
+        dt=date.today(),
+        dtt=datetime.now(),
     )
 
     schema = pydantic_to_schema(TestModel)
@@ -101,6 +120,8 @@ def test_pydantic_to_arrow_py38():
                 ),
                 False,
             ),
+            pa.field("dt", pa.date32(), False),
+            pa.field("dtt", pa.timestamp("us"), False),
         ]
     )
     assert schema == expect_schema