Repository: pyjanitor-devs/pyjanitor (python)

Issue:
callable type hint for `names_transform` doesn't work
# Brief Description

astype does not work with callable, the type hint should be for a pandas extension dtype, or a numpy dtype, or a python type

# System Information

- Operating system: macOS
- Python version (required): 3.9

# Minimally Reproducible Code

<!-- If you provide minimal code that reproduces the problem, this makes it easier for us to debug what's going on.

Minimal code should be trivially copy/pastable into a Python interpreter in its entirety. Be sure to include imports.
-->

# Error Messages

```
In [215]: df
Out[215]:
   id      start        end
0   1 2019-01-01 2019-01-03
1   1 2019-01-02 2019-01-04
2   2 2019-01-03 2019-01-05
3   2 2019-01-04 2019-01-06
4   1 2019-01-05 2019-01-07


# no changes
df.pivot_longer(index = 'id', ignore_index = False, sort_by_appearance=False, names_transform=pd.CategoricalDtype).dtypes
id                   int64
variable            object
value       datetime64[ns]
dtype: object


df.pivot_longer(index = 'id', ignore_index = False, sort_by_appearance=False, 
                            names_transform=pd.CategoricalDtype(['start', 'end'], ordered = True))

TypeError: names_transform should be one of [<class 'str'>, typing.Callable, <class 'dict'>].
```


Pull request description:
[ENH] remove check on names_transform
# PR Description

Please describe the changes proposed in the pull request:

- remove check on names_transform ... checks will be done by `pd.astype`
- faster string concatenation for multiple `.value`

**This PR resolves #1119.**

# PR Checklist

<!-- This checklist exists for newcomers who are not yet familiar with our requirements. If you are experienced with
the project, please feel free to delete this section. -->

Please ensure that you have done the following:

1. [ ] PR in from a fork off your branch. Do not PR from `<your_username>`:`dev`, but rather from `<your_username>`:`<feature-branch_name>`.
<!-- Doing this helps us keep the commit history much cleaner than it would otherwise be. -->
2. [x] If you're not on the contributors list, add yourself to `AUTHORS.md`.
<!-- We'd like to acknowledge your contributions! -->
3. [x] Add a line to `CHANGELOG.md` under the latest version header (i.e. the one that is "on deck") describing the contribution.
    - Do use some discretion here; if there are multiple PRs that are related, keep them in a single line.

# Automatic checks

There will be automatic checks run on the PR. These include:

- Building a preview of the docs on Netlify
- Automatically linting the code
- Making sure the code is documented
- Making sure that all tests are passed
- Making sure that code coverage doesn't go down.

# Relevant Reviewers

<!-- Finally, please tag relevant maintainers to review. -->

Please tag maintainers to review.

- @ericmjl


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index dd5ec74..7467a35 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,7 @@
 -   [ENH] `min_max_scale` drop `old_min` and `old_max` to fit sklearn's method API. Issue #1068 @Zeroto521
 -   [ENH] Add `jointly` option for `min_max_scale` support to transform each column values or entire values. Default transform each column, similar behavior to `sklearn.preprocessing.MinMaxScaler`. (Issue #1067, PR #1112, PR #1123) @Zeroto521
 -   [INF] Require pyspark minimal version is v3.2.0 to cut duplicates codes. Issue #1110 @Zeroto521
+-   [ENH] Added support for extension arrays in `expand_grid`. Issue #1121 @samukweku
 
 
 ## [v0.23.1] - 2022-05-03
diff --git a/examples/notebooks/expand_grid.ipynb b/examples/notebooks/expand_grid.ipynb
index e4e57e3..80dbaad 100644
--- a/examples/notebooks/expand_grid.ipynb
+++ b/examples/notebooks/expand_grid.ipynb
@@ -1108,11 +1108,9 @@
   }
  ],
  "metadata": {
-  "interpreter": {
-   "hash": "98b0a9b7b4eaaa670588a142fd0a9b87eaafe866f1db4228be72b4211d12040f"
-  },
   "kernelspec": {
-   "display_name": "Python 3.8.10 64-bit ('base': conda)",
+   "display_name": "Python 3.9.10 ('base')",
+   "language": "python",
    "name": "python3"
   },
   "language_info": {
@@ -1125,7 +1123,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.9"
+   "version": "3.9.10"
   },
   "orig_nbformat": 2
  },
diff --git a/janitor/functions/pivot.py b/janitor/functions/pivot.py
index a377d50..cecba97 100644
--- a/janitor/functions/pivot.py
+++ b/janitor/functions/pivot.py
@@ -11,7 +11,6 @@ from pandas.api.types import (
     is_list_like,
     is_string_dtype,
 )
-from pandas.core.construction import extract_array
 from pandas.core.dtypes.concat import concat_compat
 
 from janitor.functions.utils import (
@@ -258,8 +257,8 @@ def pivot_longer(
         arguments must match.
         `names_pattern` does not work with MultiIndex columns.
     :param names_transform: Use this option to change the types of columns that
-        have been transformed to rows. This does not applies to the values columns.
-        Accepts a string, a callable, or a dictionary that is acceptable by `pd.astype`.
+        have been transformed to rows. This does not applies to the values' columns.
+        Accepts any argument that is acceptable by `pd.astype`.
     :param sort_by_appearance: Default `False`. Boolean value that determines
         the final look of the DataFrame. If `True`, the unpivoted DataFrame
         will be stacked in order of first appearance.
@@ -526,9 +525,6 @@ def _data_checks_pivot_longer(
                     "index parameter. Kindly provide unique label(s)."
                 )
 
-    if names_transform:
-        check("names_transform", names_transform, [str, Callable, dict])
-
     check("sort_by_appearance", sort_by_appearance, [bool])
 
     check("ignore_index", ignore_index, [bool])
@@ -602,9 +598,7 @@ def _computations_pivot_longer(
         return df
 
     if index:
-        index = {
-            name: extract_array(df[name], extract_numpy=True) for name in index
-        }
+        index = {name: df[name]._values for name in index}
 
     if len(column_names) != len(set(column_names)):
         column_names = pd.unique(column_names)
@@ -732,10 +726,7 @@ def _pivot_longer_names_pattern_sequence(
             values = _names_transform(
                 names_transform, is_dataframe=False, values=values
             )
-        values = {
-            name: extract_array(arr, extract_numpy=True).repeat(len_index)
-            for name, arr in values
-        }
+        values = {name: arr._values.repeat(len_index) for name, arr in values}
     else:
         values = {}
 
@@ -896,12 +887,9 @@ def _base_melt(
         outcome = _names_transform(
             names_transform, is_dataframe=False, values=outcome
         )
-    outcome = {
-        name: extract_array(arr, extract_numpy=True).repeat(len_index)
-        for name, arr in outcome
-    }
+    outcome = {name: arr._values.repeat(len_index) for name, arr in outcome}
 
-    values = [extract_array(arr, extract_numpy=True) for _, arr in df.items()]
+    values = [arr._values for _, arr in df.items()]
     values = {values_to: concat_compat(values)}
 
     return _final_frame_longer(
@@ -935,8 +923,15 @@ def _pivot_longer_dot_value(
     """
     if np.count_nonzero(mapping.columns == ".value") > 1:
         outcome = mapping.pop(".value")
-        outcome = outcome.agg("".join, axis=1)
-        mapping[".value"] = outcome
+        out = outcome.iloc[:, 0]
+        # for loop preferred over agg
+        # primarily for speed
+        # if the column is a large array
+        # direct addition is surprisingly faster than
+        # the convenient agg(','.join, axis = 1) option
+        for _, val in outcome.iloc[:, 1:].items():
+            out += val
+        mapping[".value"] = out
 
     exclude = {
         word
@@ -1001,7 +996,7 @@ def _pivot_longer_dot_value(
                 names_transform, is_dataframe=True, values=outcome
             )
         outcome = {
-            name: extract_array(arr, extract_numpy=True).repeat(len_index)
+            name: arr._values.repeat(len_index)
             for name, arr in outcome.items()
         }
 
@@ -1064,8 +1059,7 @@ def _dict_from_grouped_names(df: pd.DataFrame) -> dict:
     """
     outcome = defaultdict(list)
     for num, name in enumerate(df.columns):
-        arr = df.iloc[:, num]
-        arr = extract_array(arr, extract_numpy=True)
+        arr = df.iloc[:, num]._values
         outcome[name].append(arr)
     return {name: concat_compat(arr) for name, arr in outcome.items()}
 
diff --git a/janitor/functions/utils.py b/janitor/functions/utils.py
index 99dbfff..3c156fa 100644
--- a/janitor/functions/utils.py
+++ b/janitor/functions/utils.py
@@ -188,7 +188,7 @@ def _computations_expand_grid(others: dict) -> pd.DataFrame:
     grid = ((*left, right) for left, right in grid)
     contents = {}
     for key, value, grid_index in grid:
-        contents = {**contents, **_expand_grid(value, grid_index, key)}
+        contents.update(_expand_grid(value, grid_index, key))
     return pd.DataFrame(contents, copy=False)
 
 
diff --git a/janitor/utils.py b/janitor/utils.py
index b01b979..3ee11d5 100644
--- a/janitor/utils.py
+++ b/janitor/utils.py
@@ -9,7 +9,6 @@ from typing import Callable, Dict, Iterable, Union
 
 import numpy as np
 import pandas as pd
-from pandas.core.construction import extract_array
 
 
 def check(varname: str, value, expected_types: list):
@@ -57,7 +56,6 @@ def _expand_grid(value, grid_index, key):
 def _sub_expand_grid(value, grid_index, key):  # noqa: F811
     """
     Expands the numpy array based on `grid_index`.
-
     Returns a dictionary.
     """
 
@@ -73,43 +71,37 @@ def _sub_expand_grid(value, grid_index, key):  # noqa: F811
     if value.ndim == 1:
         return {(key, 0): value}
 
-    return {(key, num): arr for num, arr in enumerate(value.T)}
+    return {(key, num): value[:, num] for num in range(value.shape[-1])}
 
 
+@_expand_grid.register(pd.api.extensions.ExtensionArray)
 @_expand_grid.register(pd.arrays.PandasArray)
 def _sub_expand_grid(value, grid_index, key):  # noqa: F811
     """
     Expands the pandas array based on `grid_index`.
-
     Returns a dictionary.
     """
 
-    value = value[grid_index]
-
-    return {(key, 0): value}
+    return {(key, 0): value[grid_index]}
 
 
+@_expand_grid.register(pd.Index)
 @_expand_grid.register(pd.Series)
 def _sub_expand_grid(value, grid_index, key):  # noqa: F811
     """
-    Expands the Series based on `grid_index`.
-
+    Expands the pd.Series/pd.Index based on `grid_index`.
     Returns a dictionary.
     """
 
-    name = value.name
-    if not name:
-        name = 0
-    value = extract_array(value, extract_numpy=True)[grid_index]
+    name = value.name or 0
 
-    return {(key, name): value}
+    return {(key, name): value._values[grid_index]}
 
 
 @_expand_grid.register(pd.DataFrame)
 def _sub_expand_grid(value, grid_index, key):  # noqa: F811
     """
     Expands the DataFrame based on `grid_index`.
-
     Returns a dictionary.
     """
 
@@ -120,8 +112,7 @@ def _sub_expand_grid(value, grid_index, key):  # noqa: F811
         value = value.set_axis(columns, axis="columns")
 
     return {
-        (key, name): extract_array(val, extract_numpy=True)[grid_index]
-        for name, val in value.items()
+        (key, name): val._values[grid_index] for name, val in value.items()
     }
 
 
@@ -129,7 +120,6 @@ def _sub_expand_grid(value, grid_index, key):  # noqa: F811
 def _sub_expand_grid(value, grid_index, key):  # noqa: F811
     """
     Expands the MultiIndex based on `grid_index`.
-
     Returns a dictionary.
     """
 
@@ -138,7 +128,7 @@ def _sub_expand_grid(value, grid_index, key):  # noqa: F811
     for n in range(value.nlevels):
         arr = value.get_level_values(n)
         name = arr.name
-        arr = extract_array(arr, extract_numpy=True)[grid_index]
+        arr = arr._values[grid_index]
         if not name:
             name = num
             num += 1
@@ -146,19 +136,6 @@ def _sub_expand_grid(value, grid_index, key):  # noqa: F811
     return contents
 
 
-@_expand_grid.register(pd.Index)
-def _sub_expand_grid(value, grid_index, key):  # noqa: F811
-    """
-    Expands the Index based on `grid_index`.
-
-    Returns a dictionary.
-    """
-    name = value.name
-    if not name:
-        name = 0
-    return {(key, name): extract_array(value, extract_numpy=True)[grid_index]}
-
-
 def import_message(
     submodule: str,
     package: str,


Test patch:
diff --git a/tests/functions/test_expand_grid.py b/tests/functions/test_expand_grid.py
index c54c73d..aacda72 100644
--- a/tests/functions/test_expand_grid.py
+++ b/tests/functions/test_expand_grid.py
@@ -8,6 +8,7 @@ from janitor.testing_utils.strategies import (
     categoricaldf_strategy,
 )
 from janitor.functions import expand_grid
+from functools import reduce
 
 
 @given(df=df_strategy())
@@ -324,3 +325,21 @@ def test_series_name(df):
         [["city", "A"], ["cities", 0]]
     )
     assert_frame_equal(result, expected)
+
+
+def test_extension_array():
+    """Test output on an extension array"""
+    others = dict(
+        id=pd.Categorical(
+            values=(2, 1, 1, 2, 1), categories=(1, 2, 3), ordered=True
+        ),
+        year=(2018, 2018, 2019, 2020, 2020),
+        gender=pd.Categorical(("female", "male", "male", "female", "male")),
+    )
+
+    expected = expand_grid(others=others).droplevel(axis=1, level=-1)
+    others = [pd.Series(val).rename(key) for key, val in others.items()]
+
+    func = lambda x, y: pd.merge(x, y, how="cross")  # noqa: E731
+    actual = reduce(func, others)
+    assert_frame_equal(expected, actual)