Repository: posit-dev/positron
Pull request #6739: Data Explorer: Support pandas.CategoricalDtype, polars.Categorical, enable sparklines for Polars
Author: login: wesm
Created: 2025-03-11T21:36:19Z | Merged: 2025-03-12T20:40:44Z

Addresses #3417 for Python. Categorical/Factor data types should return their schemas as the display type of the categories, which for R and Polars should be string, but for pandas it is whatever the display type of the dtype's categories member is.

I also noted that sparklines were disabled for Polars, but they had unit tests so I enabled them and they should work now!

e2e: @:data-explorer

### Release Notes

#### New Features

- pandas and Polars Categorical data types now display correctly in the data explorer.
- Polars sparklines are shown in the summary pane of the data explorer.

#### Bug Fixes

- N/A


### QA Notes

I used this code to create some categorical data (note that only strings can be categorical in polars):

```python
import pandas as pd
import polars as pl

# Example data frame with categorical column types
df = pd.DataFrame({'a': ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j'],
    'b': list(range(10))
    })
df['a'] = df['a'].astype('category')
df['b'] = df['b'].astype('category')

pdf = pl.from_pandas(df)
``` 


Diff:
diff --git a/extensions/positron-python/python_files/posit/positron/data_explorer.py b/extensions/positron-python/python_files/posit/positron/data_explorer.py
index bea732e4b0e1..1f70632aad9d 100644
--- a/extensions/positron-python/python_files/posit/positron/data_explorer.py
+++ b/extensions/positron-python/python_files/posit/positron/data_explorer.py
@@ -813,9 +813,9 @@ def _box_number_stats(min_val, max_val, mean_val, median_val, std_val):
     )
 
 
-def _box_other_stats(num_unique):
+def _box_other_stats(num_unique, type_display=ColumnDisplayType.Object):
     return ColumnSummaryStats(
-        type_display=ColumnDisplayType.Object,
+        type_display=type_display,
         other_stats=SummaryStatsOther(num_unique=int(num_unique)),
     )
 
@@ -1026,9 +1026,13 @@ def _pandas_summarize_string(col: pd.Series, _options: FormatOptions):
     return _box_string_stats(num_empty, num_unique)
 
 
-def _pandas_summarize_object(col: pd.Series, _options: FormatOptions):
+def _pandas_summarize_object(
+    col: pd.Series,
+    _options: FormatOptions,
+    type_display: ColumnDisplayType = ColumnDisplayType.Object,
+):
     num_unique = col.nunique()
-    return _box_other_stats(num_unique)
+    return _box_other_stats(num_unique, type_display=type_display)
 
 
 def _pandas_summarize_boolean(col: pd.Series, _options: FormatOptions):
@@ -1313,6 +1317,8 @@ def _get_inferred_dtype(cls, column, column_index: int, state: DataExplorerState
 
     @classmethod
     def _get_type(cls, column, column_index, state: DataExplorerState):
+        import pandas as pd
+
         # A helper function for returning the backend type_name and
         # the display type when returning schema results or analyzing
         # schema changes
@@ -1321,11 +1327,22 @@ def _get_type(cls, column, column_index, state: DataExplorerState):
         if dtype == object:  # noqa: E721
             type_name = cls._get_inferred_dtype(column, column_index, state)
             type_name = cls.TYPE_NAME_MAPPING.get(type_name, type_name)
+            type_display = cls._get_type_display(type_name)
+        elif isinstance(dtype, pd.CategoricalDtype):
+            type_name = str(dtype)
+            if dtype.categories.dtype == object:
+                categories_type_name = cls._get_inferred_dtype(
+                    dtype.categories, column_index, state
+                )
+                type_display = cls.TYPE_NAME_MAPPING.get(categories_type_name, categories_type_name)
+            else:
+                categories_type_name = str(dtype.categories.dtype)
+                type_display = cls._get_type_display(categories_type_name)
         else:
             # TODO: more sophisticated type mapping
             type_name = str(dtype)
+            type_display = cls._get_type_display(type_name)
 
-        type_display = cls._get_type_display(type_name)
         return type_name, type_display
 
     TYPE_DISPLAY_MAPPING = MappingProxyType(
@@ -1351,7 +1368,6 @@ def _get_type(cls, column, column_index, state: DataExplorerState):
             "mixed": "object",
             "decimal": "number",
             "complex": "number",
-            "categorical": "categorical",
             "bool": "boolean",
             "datetime64": "datetime",
             "datetime64[ns]": "datetime",
@@ -2020,6 +2036,15 @@ def _polars_summarize_string(col: pl.Series, _):
     return _box_string_stats(num_empty, num_unique)
 
 
+def _polars_summarize_object(
+    col: pl.Series,
+    _format_options: FormatOptions,
+    type_display: ColumnDisplayType = ColumnDisplayType.Object,
+):
+    num_unique = col.n_unique()
+    return _box_other_stats(num_unique, type_display=type_display)
+
+
 def _polars_summarize_boolean(col: pl.Series, _):
     null_count = col.is_null().sum()
     true_count = col.sum()
@@ -2199,12 +2224,22 @@ def _construct_schema(
         column_name: str,
         column_index: int,
     ):
-        type_display = cls._get_type_display(column.dtype)
+        import polars as pl
+
+        if isinstance(column.dtype, pl.Categorical):
+            # Categorical is always string in polars
+            type_display = "string"
+            # For Categorical types, we just use "Categorical" for the type name
+            # for simplicity
+            type_name = "Categorical"
+        else:
+            type_display = cls._get_type_display(column.dtype)
+            type_name = str(column.dtype)
 
         return ColumnSchema(
             column_name=column_name,
             column_index=column_index,
-            type_name=str(column.dtype),
+            type_name=type_name,
             type_display=type_display,
         )
 
@@ -2230,7 +2265,7 @@ def _construct_schema(
             "Object": "object",
             "List": "array",
             "Struct": "struct",
-            "Categorical": "unknown",  # See #3417
+            "Categorical": "categorical",
             "Duration": "unknown",  # See #3418
             "Enum": "unknown",
             "Null": "unknown",  # Not yet implemented
@@ -2547,6 +2582,7 @@ def _get_column(self, column_index: int) -> pl.Series:
             ColumnDisplayType.Boolean: _polars_summarize_boolean,
             ColumnDisplayType.Number: _polars_summarize_number,
             ColumnDisplayType.String: _polars_summarize_string,
+            ColumnDisplayType.Object: _polars_summarize_object,
             ColumnDisplayType.Date: _polars_summarize_date,
             ColumnDisplayType.Datetime: _polars_summarize_datetime,
         }
@@ -2643,6 +2679,22 @@ def _prof_histogram(
                     profile_type=ColumnProfileType.SummaryStats,
                     support_status=SupportStatus.Supported,
                 ),
+                ColumnProfileTypeSupportStatus(
+                    profile_type=ColumnProfileType.SmallHistogram,
+                    support_status=SupportStatus.Supported,
+                ),
+                ColumnProfileTypeSupportStatus(
+                    profile_type=ColumnProfileType.LargeHistogram,
+                    support_status=SupportStatus.Supported,
+                ),
+                ColumnProfileTypeSupportStatus(
+                    profile_type=ColumnProfileType.SmallFrequencyTable,
+                    support_status=SupportStatus.Supported,
+                ),
+                ColumnProfileTypeSupportStatus(
+                    profile_type=ColumnProfileType.LargeFrequencyTable,
+                    support_status=SupportStatus.Supported,
+                ),
             ],
         ),
         export_data_selection=ExportDataSelectionFeatures(
diff --git a/extensions/positron-python/python_files/posit/positron/tests/test_data_explorer.py b/extensions/positron-python/python_files/posit/positron/tests/test_data_explorer.py
index c894c494b3e1..9930cc844a56 100644
--- a/extensions/positron-python/python_files/posit/positron/tests/test_data_explorer.py
+++ b/extensions/positron-python/python_files/posit/positron/tests/test_data_explorer.py
@@ -744,6 +744,17 @@ def test_pandas_get_schema(dxf: DataExplorerFixture):
             "datetime64[ns, US/Eastern]",
             "datetime",
         ),
+        # categorical
+        (
+            pd.Series(["foo", "bar", "foo", "baz", "qux"], dtype="category"),
+            "category",
+            "string",
+        ),
+        (
+            pd.Series([0, 1, 0, 1, 0], dtype="category"),
+            "category",
+            "number",
+        ),
     ]
 
     if hasattr(np, "complex256"):
@@ -3066,7 +3077,8 @@ def test_profile_histogram_windows_int32_bug():
         "Struct({'a': Int64, 'b': List(String)})",
         "struct",
     ),
-    # (pl.Object, ["Hello", True, None, 5], "Object", "object"),
+    (pl.Categorical, ["a", "b", "a", None], "Categorical", "string"),
+    (pl.Object, ["Hello", True, None, 5], "Object", "object"),
 ]
 
 
@@ -3152,6 +3164,22 @@ def test_polars_get_state(dxf: DataExplorerFixture):
             profile_type="summary_stats",
             support_status=SupportStatus.Supported,
         ),
+        ColumnProfileTypeSupportStatus(
+            profile_type="small_histogram",
+            support_status=SupportStatus.Supported,
+        ),
+        ColumnProfileTypeSupportStatus(
+            profile_type="large_histogram",
+            support_status=SupportStatus.Supported,
+        ),
+        ColumnProfileTypeSupportStatus(
+            profile_type="small_frequency_table",
+            support_status=SupportStatus.Supported,
+        ),
+        ColumnProfileTypeSupportStatus(
+            profile_type="large_frequency_table",
+            support_status=SupportStatus.Supported,
+        ),
     ]
 
 
@@ -3200,7 +3228,8 @@ def test_polars_get_data_values(dxf: DataExplorerFixture):
             _VALUE_NULL,
             "{'a': 0, 'b': None}",
         ],  # Struct({'a': Int64, 'b': List(String)}),
-        # ["Hello", "True", _VALUE_NULL, "5"],  # Object
+        ["a", "b", "a", _VALUE_NULL],  # Categorical
+        ["Hello", "True", _VALUE_NULL, "5"],  # Object
     ]
 
     assert result["columns"] == expected_columns


Reviews:
user: login: austin3dickey
body: LGTM, just some nits
state: DISMISSED
author_association: CONTRIBUTOR
submitted_at: 2025-03-12T14:53:19Z
user: login: wesm
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-12T18:50:04Z
user: login: wesm
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-12T18:50:32Z

Review comments:
diff_hunk: @@ -1321,11 +1327,22 @@ def _get_type(cls, column, column_index, state: DataExplorerState):
         if dtype == object:  # noqa: E721
             type_name = cls._get_inferred_dtype(column, column_index, state)
             type_name = cls.TYPE_NAME_MAPPING.get(type_name, type_name)
+            type_display = cls._get_type_display(type_name)
+        elif isinstance(dtype, pd.CategoricalDtype):
+            type_name = "category"
path: extensions/positron-python/python_files/posit/positron/data_explorer.py
user: login: austin3dickey
body: nit: should this be `categorical` to match the pandas name of the dtype? I do like brevity though
created_at: 2025-03-12T14:42:26Z
author_association: CONTRIBUTOR
path: extensions/positron-python/python_files/posit/positron/data_explorer.py
user: login: austin3dickey
body: When testing this locally I added another column with ordered custom categories:
```py
pd.Categorical(
    ['a', 'b', 'c', 'a'] * 2 + [np.nan] * 2, 
    ordered=True, 
    categories=['c', 'a', 'b']
)
```
Looks like when you "sort ascending" it just works! 
<img width="126" alt="image" src="https://github.com/user-attachments/assets/fd8d46ff-8b95-4367-b9bf-ea712f9d81fd" />

Though I noticed the sparkline looks like it's sorted by the string values of the categories, not the underlying custom order:
<img width="278" alt="image" src="https://github.com/user-attachments/assets/0209c012-2479-485d-b746-4de26600ec60" />

This doesn't seem like a big deal so it could be future work.
created_at: 2025-03-12T14:51:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1321,11 +1327,22 @@ def _get_type(cls, column, column_index, state: DataExplorerState):
         if dtype == object:  # noqa: E721
             type_name = cls._get_inferred_dtype(column, column_index, state)
             type_name = cls.TYPE_NAME_MAPPING.get(type_name, type_name)
+            type_display = cls._get_type_display(type_name)
+        elif isinstance(dtype, pd.CategoricalDtype):
+            type_name = "category"
path: extensions/positron-python/python_files/posit/positron/data_explorer.py
user: login: wesm
body: "category" is what `str(dtype)` returns so I will replace this hard-coded string with that for clarity. 
created_at: 2025-03-12T18:50:04Z
author_association: CONTRIBUTOR
path: extensions/positron-python/python_files/posit/positron/data_explorer.py
user: login: wesm
body: Yeah, we can improve that, especially once we have tooltips for the sparklines
created_at: 2025-03-12T18:50:32Z
author_association: CONTRIBUTOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-03-11T21:36:38Z
body: <!-- PR Tags -->
<!-- 
🚨 RED ALERT! ✋ Rule breaker detected! Tags don’t go here, they belong above ^ in the PR description using this proper format: `@:tag`. Changing them here won't do anything (trust us, we’ve tried). Confused? Check out the README hyperlink.
-->

**E2E Tests** 🚀
This PR will run tests tagged with: `@:critical` `@:data-explorer`

<sup>[readme](https://github.com/posit-dev/positron/blob/main/test/e2e/README.md#pull-requests-and-test-tags)</sup>&nbsp;&nbsp;<sup>[valid tags](https://github.com/posit-dev/positron/blob/main/test/e2e/infra/test-runner/test-tags.ts)</sup>
author_association: NONE
user: login: wesm
created_at: 2025-03-11T21:37:17Z
body: @dfalbel would you or someone else be interested in adding unit tests for factor data types on the R side of this in Ark?
author_association: CONTRIBUTOR