feat(python): Minor type-inference update for read_database (pola-r…

…s#15809)
rancomp · Apr 21, 2024 · 53c8ec9 · 53c8ec9
1 parent e3d0ef2
commit 53c8ec9
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/py-polars/polars/io/database/_inference.py b/py-polars/polars/io/database/_inference.py
@@ -22,6 +22,7 @@
     Int32,
     Int64,
     List,
+    Null,
     String,
     Time,
     UInt8,
@@ -164,6 +165,10 @@ def _infer_dtype_from_database_typename(
     elif value.startswith("BOOL"):
         dtype = Boolean
 
+    # null dtype; odd, but valid
+    elif value == "NULL":
+        dtype = Null
+
     # temporal dtypes
     elif value.startswith(("DATETIME", "TIMESTAMP")) and not (value.endswith("[D]")):
         if any((tz in value.replace(" ", "")) for tz in ("TZ", "TIMEZONE")):
@@ -173,7 +178,7 @@ def _infer_dtype_from_database_typename(
         dtype = Datetime(time_unit=(unit or "us"))  # type: ignore[arg-type]
     else:
         value = re.sub(r"\d", "", value)
-        if value in ("INTERVAL", "TIMEDELTA"):
+        if value in ("INTERVAL", "TIMEDELTA", "DURATION"):
             dtype = Duration
         elif value == "DATE":
             dtype = Date

diff --git a/py-polars/tests/unit/io/database/test_inference.py b/py-polars/tests/unit/io/database/test_inference.py
@@ -58,13 +58,17 @@
         ("timestamp(5)", pl.Datetime("us")),
         ("timestamp(7)", pl.Datetime("ns")),
         ("datetime without tz", pl.Datetime("us")),
+        ("duration(2)", pl.Duration("ms")),
+        ("interval", pl.Duration("us")),
         ("date", pl.Date),
         ("time", pl.Time),
         ("date32", pl.Date),
         ("time64", pl.Time),
         # binary types
         ("BYTEA", pl.Binary),
         ("BLOB", pl.Binary),
+        # miscellaneous
+        ("NULL", pl.Null),
     ],
 )
 def test_dtype_inference_from_string(