Repository: narwhals-dev/narwhals
Pull request #382 (2024-07-02 13:11:39)

feat: add .cat.get_categories to pyarrow
<!--
# Thanks for contributing a pull request! 
## Please make sure you see our contribution guidelines: https://github.com/narwhals-dev/narwhals/blob/main/CONTRIBUTING.md
-->

## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [ ] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue # 
- Closes #

## Checklist

- [ ] Code follows style guide (ruff)
- [ ] Tests added 
- [ ] Documented the changes

## If you have comments or can explain your changes, please do so below.




Patch:
diff --git a/README.md b/README.md
index 1ee3056e..7481cfd9 100644
--- a/README.md
+++ b/README.md
@@ -14,6 +14,7 @@ Extremely lightweight and extensible compatibility layer between Polars, pandas,
 - [Read the documentation](https://narwhals-dev.github.io/narwhals/)
 - [Chat with us on Discord!](https://discord.gg/V3PqtB4VA4)
 - [Join our community call](https://calendar.google.com/calendar/embed?src=27ff6dc5f598c1d94c1f6e627a1aaae680e2fac88f848bda1f2c7946ae74d5ab%40group.calendar.google.com&ctz=Europe%2FLondon).
+- [Read the contributing guide](https://github.com/narwhals-dev/narwhals/blob/main/CONTRIBUTING.md).
 
 Seamlessly support all, without depending on any!
 
diff --git a/narwhals/_arrow/series.py b/narwhals/_arrow/series.py
index 4fdafb97..84725f98 100644
--- a/narwhals/_arrow/series.py
+++ b/narwhals/_arrow/series.py
@@ -6,6 +6,7 @@ from typing import Iterable
 
 from narwhals._arrow.utils import translate_dtype
 from narwhals._pandas_like.utils import native_series_from_iterable
+from narwhals.dependencies import get_pyarrow
 from narwhals.dependencies import get_pyarrow_compute
 
 if TYPE_CHECKING:
@@ -89,6 +90,10 @@ class ArrowSeries:
     def dt(self) -> ArrowSeriesDateTimeNamespace:
         return ArrowSeriesDateTimeNamespace(self)
 
+    @property
+    def cat(self) -> ArrowSeriesCatNamespace:
+        return ArrowSeriesCatNamespace(self)
+
 
 class ArrowSeriesDateTimeNamespace:
     def __init__(self, series: ArrowSeries) -> None:
@@ -101,3 +106,17 @@ class ArrowSeriesDateTimeNamespace:
         # https://arrow.apache.org/docs/python/generated/pyarrow.compute.strftime.html
         format = format.replace("%S.%f", "%S").replace("%S%.f", "%S")
         return self._series._from_series(pc.strftime(self._series._series, format))
+
+
+class ArrowSeriesCatNamespace:
+    def __init__(self, series: ArrowSeries) -> None:
+        self._series = series
+
+    def get_categories(self) -> ArrowSeries:
+        pa = get_pyarrow()
+        ca = self._series._series
+        # todo: this looks potentially expensive - is there no better way?
+        out = pa.chunked_array(
+            [pa.concat_arrays([x.dictionary for x in ca.chunks]).unique()]
+        )
+        return self._series._from_series(out)


Test patch:
diff --git a/tests/expr/cat/get_categories_test.py b/tests/expr/cat/get_categories_test.py
index fabce934..66088173 100644
--- a/tests/expr/cat/get_categories_test.py
+++ b/tests/expr/cat/get_categories_test.py
@@ -2,6 +2,8 @@ from __future__ import annotations
 
 from typing import Any
 
+import pyarrow as pa
+
 import narwhals as nw
 from tests.utils import compare_dicts
 
@@ -16,3 +18,15 @@ def test_get_categories(constructor: Any) -> None:
     compare_dicts(result, expected)
     result = df.select(df["a"].cat.get_categories())
     compare_dicts(result, expected)
+
+
+def test_get_categories_pyarrow() -> None:
+    # temporary test until we have `cast` in pyarrow - later, fuse
+    # this with test above
+    table = pa.table(
+        {"a": pa.array(["a", "b", None, "d"], pa.dictionary(pa.int64(), pa.utf8()))}
+    )
+    df = nw.from_native(table, eager_only=True)
+    result = df["a"].cat.get_categories().to_list()
+    expected = ["a", "b", "d"]
+    assert result == expected