fix(rust, python): fix logical types in arr.get (#7094)

pola-rs · Feb 22, 2023 · 5c3b599 · 5c3b599
1 parent be61a93
commit 5c3b599
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git a/polars/polars-ops/src/chunked_array/list/namespace.rs b/polars/polars-ops/src/chunked_array/list/namespace.rs
@@ -316,6 +316,8 @@ pub trait ListNameSpaceImpl: AsList {
             .map(|arr| sublist_get(arr, idx))
             .collect::<Vec<_>>();
         Series::try_from((ca.name(), chunks))
+            .unwrap()
+            .cast(&ca.inner_dtype())
     }
 
     #[cfg(feature = "list_take")]

diff --git a/py-polars/tests/unit/namespaces/test_list.py b/py-polars/tests/unit/namespaces/test_list.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import typing
+from datetime import date, datetime
 
 import numpy as np
 import pytest
@@ -377,3 +378,17 @@ def test_list_function_group_awareness() -> None:
         "take": [[[100]], [[105]], [[100]]],
         "slice": [[[100, 103]], [[105, 106, 105]], [[100, 102]]],
     }
+
+
+def test_list_get_logical_types() -> None:
+    df = pl.DataFrame(
+        {
+            "date_col": [[datetime(2023, 2, 1).date(), datetime(2023, 2, 2).date()]],
+            "datetime_col": [[datetime(2023, 2, 1), datetime(2023, 2, 2)]],
+        }
+    )
+
+    assert df.select(pl.all().arr.get(1).suffix("_element_1")).to_dict(False) == {
+        "date_col_element_1": [date(2023, 2, 2)],
+        "datetime_col_element_1": [datetime(2023, 2, 2, 0, 0)],
+    }