fix: Convert Polars container dtype entries to list in _get_cell (#778) (#794)

* fix: convert container type entries to list in polars tables

* refactor: use tuple for second arg in isinstance

* test: test _get_cell for container dtypes

* Guard None before to_list() for list cells

* Handle None in container dtype tests

* Include pl.Array in list column detection

* refactor: simplify list-column check

---------

Co-authored-by: Michael Chow <mc_al_github@fastmail.com>
Co-authored-by: Richard Iannone <riannone@me.com>

Author: 3 people

great_tables/_tbl_data.py

@@ -224,7 +224,15 @@ def _get_cell(data: DataFrameLike, row: int, column: str) -> Any:
 
 @_get_cell.register(PlDataFrame)
 def _(data: Any, row: int, column: str) -> Any:
-    return data[column][row]
+    import polars as pl
+
+    res = data[column][row]
+
+    # container dtypes (pl.List, pl.Array) return a pl.Series
+    if isinstance(res, pl.Series):
+        return res.to_list()
+
+    return res
 
 
 @_get_cell.register(PdDataFrame)
@@ -570,7 +578,9 @@ def _(df: PlDataFrame):
     import polars.selectors as cs
 
     list_cols = [
-        name for name, dtype in df.schema.items() if issubclass(dtype.base_type(), pl.List)
+        name
+        for name, dtype in df.schema.items()
+        if issubclass(dtype.base_type(), (pl.List, pl.Array))
     ]
 
     return df.with_columns(

tests/test_tbl_data.py

@@ -1,11 +1,12 @@
 import math
+
 import pandas as pd
 import polars as pl
-import pyarrow as pa
 import polars.testing
+import pyarrow as pa
 import pytest
+
 from great_tables import GT
-from great_tables._utils_render_html import create_body_component_h
 from great_tables._tbl_data import (
     DataFrameLike,
     SeriesLike,
@@ -14,6 +15,7 @@
     _set_cell,
     _validate_selector_list,
     cast_frame_to_string,
+    copy_frame,
     create_empty_frame,
     eval_aggregate,
     eval_select,
@@ -24,8 +26,8 @@
     to_frame,
     to_list,
     validate_frame,
-    copy_frame,
 )
+from great_tables._utils_render_html import create_body_component_h
 
 params_frames = [
     pytest.param(pd.DataFrame, id="pandas"),
@@ -38,13 +40,37 @@
     pytest.param(pa.array, id="arrow"),
     pytest.param(lambda a: pa.chunked_array([a]), id="arrow-chunked"),
 ]
+params_pl_container_dtypes = [
+    pytest.param(pl.List, id="list"),
+    pytest.param(pl.Array, id="array"),
+]
 
 
 @pytest.fixture(params=params_frames, scope="function")
 def df(request) -> pd.DataFrame:
     return request.param({"col1": [1, 2, 3], "col2": ["a", "b", "c"], "col3": [4.0, 5.0, 6.0]})
 
 
+@pytest.fixture(params=params_pl_container_dtypes, scope="function")
+def df_container_dtypes(request):
+    dtype_constructor = request.param
+
+    if dtype_constructor == pl.List:
+        return pl.DataFrame(
+            {"col1": [1, 2, 3], "col2": [[1, 2, 3], [4, 5], None], "col3": ["a", "b", "c"]}
+        )
+    # return a pl df with pl.Array columns
+    else:
+        col2_as_array = pl.col("col2").cast(pl.Array(pl.Int32, shape=(3,)))
+        return pl.DataFrame(
+            {
+                "col1": [1, 2, 3],
+                "col2": [[1, 2, 3], [4, 5, 6], None],
+                "col3": ["a", "b", "c"],
+            }
+        ).with_columns(col2_as_array)
+
+
 @pytest.fixture(params=params_series, scope="function")
 def ser(request) -> SeriesLike:
     return request.param([1.0, 2.0, None])
@@ -75,6 +101,12 @@ def test_get_cell(df: DataFrameLike):
     assert _get_cell(df, 1, "col2") == "b"
 
 
+def test_get_cell_container_dtypes(df_container_dtypes: pl.DataFrame):
+    "Checks that container dtype entries in polars dfs are returned as lists"
+    assert isinstance(_get_cell(df_container_dtypes, 0, "col2"), list)
+    assert _get_cell(df_container_dtypes, 2, "col2") is None
+
+
 def test_set_cell(df: DataFrameLike):
     expected_data = {"col1": [1, 2, 3], "col2": ["a", "x", "c"], "col3": [4.0, 5.0, 6.0]}
     if isinstance(df, pa.Table):