Repository: modin-project/modin
Pull request: FIX-#4859: Add support for PyArrow Dictionary Arrays to type mapping
Signed-off-by: Andrey Pavlenko <andrey.a.pavlenko@gmail.com>

Fixed the issue https://github.com/modin-project/modin/pull/4864#discussion_r1030784639

- [x] first commit message and PR title follow format outlined [here](https://modin.readthedocs.io/en/latest/development/contributing.html#commit-message-formatting)
  > **_NOTE:_**  If you edit the PR title to match this format, you need to add another commit (even if it's empty) or amend your last commit for the CI job that checks the PR title to pick up the new PR title.
- [x] passes `flake8 modin/ asv_bench/benchmarks scripts/doc_checker.py`
- [x] passes `black --check modin/ asv_bench/benchmarks scripts/doc_checker.py`
- [x] signed commit with `git commit -s` <!-- you can amend your commit with a signature via `git commit -amend -s` -->
- [x] Resolves #4859 <!-- issue must be created for each patch -->
- [x] tests added and passing
- [ ] module layout described at `docs/development/architecture.rst` is up-to-date <!-- if you have added, renamed or removed files or directories please update the documentation accordingly -->


Problem statement:
Unable to create DataFrame form a PyArrow table containing a Dictionary Array
### System information
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Ubuntu 22.04.1
- **Modin version** (`modin.__version__`): 0.15.0+89.g8f3f3899
- **Python version**: 3.8.13
- **Code we can use to reproduce**:

import pyarrow as pa
import modin.config as cfg
from modin.pandas.utils import from_arrow
from modin.pandas.test.utils import df_equals

cfg.StorageFormat.put('omnisci')
cfg.IsExperimental.put(True)

indices = pa.array([0, 1, 0, 1, 2, 0, None, 2])
dictionary = pa.array(['first', 'second', 'third'])
dict_array = pa.DictionaryArray.from_arrays(indices, dictionary)
at = pa.table({"col": dict_array})
pdf = at.to_pandas()
print(pdf)
mdf = from_arrow(at)
print(mdf)
df_equals(mdf, pdf)

### Describe the problem
The method from_arrow() fails with the bellow error if the input table contains a dictionary. 

### Source code / logs
Traceback (most recent call last):
...
  File "pyarrow/types.pxi", line 200, in pyarrow.lib.DataType.to_pandas_dtype
NotImplementedError: dictionary<values=string, indices=int64, ordered=0>

Hints from issue discussion:
Reopening the issue because of https://github.com/modin-project/modin/pull/4864#discussion_r1030784639.

Commit message: FIX-#4859: Added support for PyArrow Dictionary Arrays to type mapping

Make sure dtype is not materialized after str().

Signed-off-by: Andrey Pavlenko <andrey.a.pavlenko@gmail.com>

Patch to review:
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 82c58eaafed..b902c0fb9e7 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -435,6 +435,7 @@ jobs:
       - run: pytest modin/experimental/core/execution/native/implementations/hdk_on_native/test/test_init.py
       - run: pytest modin/experimental/core/execution/native/implementations/hdk_on_native/test/test_dataframe.py
       - run: pytest modin/pandas/test/test_io.py::TestCsv --verbose
+      - run: pytest modin/pandas/test/dataframe/test_udf.py
       - run: pytest modin/test/interchange/dataframe_protocol/test_general.py
       - run: pytest modin/test/interchange/dataframe_protocol/hdk
       - run: python examples/docker/modin-hdk/census-hdk.py examples/data/census_1k.csv -no-ml
diff --git a/modin/core/dataframe/pandas/dataframe/dataframe.py b/modin/core/dataframe/pandas/dataframe/dataframe.py
index 3f31942162f..6b7292a46a9 100644
--- a/modin/core/dataframe/pandas/dataframe/dataframe.py
+++ b/modin/core/dataframe/pandas/dataframe/dataframe.py
@@ -3095,8 +3095,6 @@ def _arrow_type_to_dtype(cls, arrow_type):
         except NotImplementedError:
             if pyarrow.types.is_time(arrow_type):
                 res = np.dtype(datetime.time)
-            elif pyarrow.types.is_dictionary(arrow_type):
-                res = pandas.CategoricalDtype
             else:
                 raise
 
diff --git a/modin/experimental/core/execution/native/implementations/hdk_on_native/dataframe/dataframe.py b/modin/experimental/core/execution/native/implementations/hdk_on_native/dataframe/dataframe.py
index 86d632739d3..41583e7c014 100644
--- a/modin/experimental/core/execution/native/implementations/hdk_on_native/dataframe/dataframe.py
+++ b/modin/experimental/core/execution/native/implementations/hdk_on_native/dataframe/dataframe.py
@@ -12,6 +12,7 @@
 # governing permissions and limitations under the License.
 
 """Module provides ``HdkOnNativeDataframe`` class implementing lazy frame."""
+import pandas
 
 from modin.core.dataframe.pandas.dataframe.dataframe import PandasDataframe
 from modin.core.dataframe.base.dataframe.utils import Axis, JoinType
@@ -185,6 +186,7 @@ def __init__(
         assert len(dtypes) == len(
             self._table_cols
         ), f"unaligned dtypes ({dtypes}) and table columns ({self._table_cols})"
+
         if isinstance(dtypes, list):
             if self._index_cols is not None:
                 # Table stores both index and data columns but those are accessed
@@ -214,12 +216,17 @@ def __init__(
         if self._has_arrow_table() and self._partitions.size > 0:
             assert self._partitions.size == 1
             table = self._partitions[0][0].get()
-            if len(table) > 0 and table.column_names[0] != f"F_{self._table_cols[0]}":
-                new_names = [f"F_{col}" for col in table.column_names]
-                new_table = table.rename_columns(new_names)
+            column_names = table.column_names
+            if len(table) > 0 and column_names[0] != f"F_{self._table_cols[0]}":
+                column_names = [f"F_{col}" for col in column_names]
+                table = table.rename_columns(column_names)
                 self._partitions[0][
                     0
-                ] = self._partition_mgr_cls._partition_class.put_arrow(new_table)
+                ] = self._partition_mgr_cls._partition_class.put_arrow(table)
+
+            for i, t in enumerate(dtypes):
+                if isinstance(t, LazyProxyCategoricalDtype):
+                    dtypes[i] = t._new(table, column_names[i])
 
         self._uses_rowid = uses_rowid
         self._force_execution_mode = force_execution_mode
@@ -2508,10 +2515,13 @@ def from_arrow(cls, at, index_cols=None, index=None, columns=None):
             new_columns = pd.Index(data=at.column_names, dtype="O")
             new_index = pd.RangeIndex(at.num_rows)
 
-        new_dtypes = pd.Series(
-            [cls._arrow_type_to_dtype(col.type) for col in at.columns],
-            index=at.column_names,
-        )
+        new_dtypes = []
+
+        for col in at.columns:
+            if pyarrow.types.is_dictionary(col.type):
+                new_dtypes.append(LazyProxyCategoricalDtype(at, col._name))
+            else:
+                new_dtypes.append(cls._arrow_type_to_dtype(col.type))
 
         if len(unsupported_cols) > 0:
             ErrorMessage.single_warning(
@@ -2525,7 +2535,7 @@ def from_arrow(cls, at, index_cols=None, index=None, columns=None):
             columns=new_columns,
             row_lengths=new_lengths,
             column_widths=new_widths,
-            dtypes=new_dtypes,
+            dtypes=pd.Series(data=new_dtypes, index=at.column_names),
             index_cols=index_cols,
             has_unsupported_data=len(unsupported_cols) > 0,
         )
@@ -2556,3 +2566,58 @@ def _is_trivial_index(cls, index):
             and index.min() == 0
             and index.max() == len(index) - 1
         )
+
+
+class LazyProxyCategoricalDtype(pandas.CategoricalDtype):
+    """
+    Proxy class for lazily retrieving categorical dtypes from arrow tables.
+
+    Parameters
+    ----------
+    table : pyarrow.Table
+        Source table.
+    name : str
+        Column name.
+    """
+
+    def __init__(self, table: pyarrow.Table, name: str):
+        self._table = table
+        self._name = name
+        self._ordered = False
+        self._lazy_categories = None
+
+    def _new(self, table: pyarrow.Table, name: str) -> pandas.CategoricalDtype:
+        """
+        Create a new proxy, if either table or column name are different.
+
+        Parameters
+        ----------
+        table : pyarrow.Table
+            Source table.
+        name : str
+            Column name.
+
+        Returns
+        -------
+        pandas.CategoricalDtype
+        """
+        if self._table is None:
+            # The table has been materialized, we don't need a proxy anymore.
+            return pandas.CategoricalDtype(self.categories)
+        elif table is self._table and name == self._name:
+            return self
+        else:
+            return LazyProxyCategoricalDtype(table, name)
+
+    @property
+    def _categories(self):  # noqa: GL08
+        if self._table is not None:
+            chunks = self._table.column(self._name).chunks
+            cat = pandas.concat([chunk.dictionary.to_pandas() for chunk in chunks])
+            self._lazy_categories = self.validate_categories(cat.unique())
+            self._table = None  # The table is not required any more
+        return self._lazy_categories
+
+    @_categories.setter
+    def _set_categories(self, categories):  # noqa: GL08
+        self._lazy_categories = categories
diff --git a/modin/experimental/core/execution/native/implementations/hdk_on_native/test/test_dataframe.py b/modin/experimental/core/execution/native/implementations/hdk_on_native/test/test_dataframe.py
index 21999c86d28..dfd12cd4d8e 100644
--- a/modin/experimental/core/execution/native/implementations/hdk_on_native/test/test_dataframe.py
+++ b/modin/experimental/core/execution/native/implementations/hdk_on_native/test/test_dataframe.py
@@ -25,6 +25,7 @@
     random_state,
     test_data,
 )
+from modin.test.interchange.dataframe_protocol.hdk.utils import split_df_into_chunks
 from .utils import eval_io, ForceHdkImport, set_execution_mode, run_and_compare
 from pandas.core.dtypes.common import is_list_like
 
@@ -2393,11 +2394,41 @@ def test_dict(self):
         indices = pyarrow.array([0, 1, 0, 1, 2, 0, None, 2])
         dictionary = pyarrow.array(["first", "second", "third"])
         dict_array = pyarrow.DictionaryArray.from_arrays(indices, dictionary)
-        at = pyarrow.table({"col": dict_array})
+        at = pyarrow.table(
+            {"col1": dict_array, "col2": [1, 2, 3, 4, 5, 6, 7, 8], "col3": dict_array}
+        )
         pdf = at.to_pandas()
+        nchunks = 3
+        chunks = split_df_into_chunks(pdf, nchunks)
+        at = pyarrow.concat_tables([pyarrow.Table.from_pandas(c) for c in chunks])
         mdf = from_arrow(at)
+        assert (
+            len(
+                mdf._query_compiler._modin_frame._partitions[0][0]
+                .get()
+                .column(0)
+                .chunks
+            )
+            == nchunks
+        )
         df_equals(mdf, pdf)
 
+        mdt = mdf.dtypes[0]
+        pdt = pdf.dtypes[0]
+        assert mdt == "category"
+        assert isinstance(mdt, pandas.CategoricalDtype)
+        assert pandas.api.types.is_categorical_dtype(mdt)
+        assert str(mdt) == str(pdt)
+
+        if type(mdt) != pandas.CategoricalDtype:
+            # This is a lazy proxy.
+            # Make sure the table is not materialized yet.
+            assert mdt._table is not None
+
+        assert mdt == pdt
+        assert pdt == mdt
+        assert repr(mdt) == repr(pdt)
+
 
 if __name__ == "__main__":
     pytest.main(["-v", __file__])
diff --git a/modin/pandas/test/dataframe/test_udf.py b/modin/pandas/test/dataframe/test_udf.py
index 32a61e2b161..fe0c404e95d 100644
--- a/modin/pandas/test/dataframe/test_udf.py
+++ b/modin/pandas/test/dataframe/test_udf.py
@@ -57,6 +57,13 @@
 pytestmark = pytest.mark.filterwarnings(default_to_pandas_ignore_string)
 
 
+# Initialize the environment. If not initialized here,
+# StorageFormat.get() == "Pandas" is always true for
+# @pytest.mark.xfail() and the test test_apply_key_error
+# fails in case of HDK backend.
+pd.DataFrame()
+
+
 def test_agg_dict():
     md_df, pd_df = create_test_dfs(test_data_values[0])
     agg_dict = {pd_df.columns[0]: "sum", pd_df.columns[-1]: ("sum", "count")}
diff --git a/requirements/env_hdk.yml b/requirements/env_hdk.yml
index b8be4a82fd9..65f9eeeead3 100644
--- a/requirements/env_hdk.yml
+++ b/requirements/env_hdk.yml
@@ -21,6 +21,7 @@ dependencies:
   - scipy
   - xgboost>=1.7.1,<2.0.0
   - scikit-learn-intelex
+  - matplotlib
   - pip:
       - black
       # TODO: remove when flake8 5.x stabilizes and appears in both pip and conda-forge; see GH-#4745


Code review comments:
diff_hunk: @@ -435,6 +435,7 @@ jobs:
       - run: pytest modin/experimental/core/execution/native/implementations/hdk_on_native/test/test_init.py
       - run: pytest modin/experimental/core/execution/native/implementations/hdk_on_native/test/test_dataframe.py
       - run: pytest modin/pandas/test/test_io.py::TestCsv --verbose
+      - run: pytest modin/pandas/test/dataframe/test_udf.py
path: .github/workflows/ci.yml
text: @user1:
why is this needed?

@author:
I've found, that after applying this fix, the test test_apply_key_error xfailed. I've fixed the failure and added the test to ci. Do you think it's not required?

@user1:
The mentioned test verifies that an error raised by `DataFrame.apply()` in case of invalid arguments matches pandas. I wouldn't say we support `.apply()` for any non-pandas storage format, thus the test only has failure condition for pandas
https://github.com/modin-project/modin/blob/17e55993c952eb31fd453d98ead94a014f549115/modin/pandas/test/dataframe/test_udf.py#L120
The `.xfail` should be unconditional here, but since these tests are not intended to run on anything different from pandas storage format the existing form of `xfail` is ok.

Summarizing, I believe that the change is unnecessary in this PR.

diff_hunk: @@ -2556,3 +2566,58 @@ def _is_trivial_index(cls, index):
             and index.min() == 0
             and index.max() == len(index) - 1
         )
+
+
+class LazyProxyCategoricalDtype(pandas.CategoricalDtype):
+    """
+    Proxy class for lazily retrieving categorical dtypes from arrow tables.
+
+    Parameters
+    ----------
+    table : pyarrow.Table
+        Source table.
+    name : str
+        Column name.
+    """
+
+    def __init__(self, table: pyarrow.Table, name: str):
+        self._table = table
+        self._name = name
+        self._ordered = False
+        self._lazy_categories = None
+
+    def _new(self, table: pyarrow.Table, name: str) -> pandas.CategoricalDtype:
+        """
+        Create a new proxy, if either table or column name are different.
+
+        Parameters
+        ----------
+        table : pyarrow.Table
+            Source table.
+        name : str
+            Column name.
+
+        Returns
+        -------
+        pandas.CategoricalDtype
path: modin/experimental/core/execution/native/implementations/hdk_on_native/dataframe/dataframe.py
text: @user1:
```suggestion
        pandas.CategoricalDtype or LazyProxyCategoricalDtype
```

diff_hunk: @@ -2393,11 +2394,41 @@ def test_dict(self):
         indices = pyarrow.array([0, 1, 0, 1, 2, 0, None, 2])
         dictionary = pyarrow.array(["first", "second", "third"])
         dict_array = pyarrow.DictionaryArray.from_arrays(indices, dictionary)
-        at = pyarrow.table({"col": dict_array})
+        at = pyarrow.table(
+            {"col1": dict_array, "col2": [1, 2, 3, 4, 5, 6, 7, 8], "col3": dict_array}
+        )
         pdf = at.to_pandas()
+        nchunks = 3
+        chunks = split_df_into_chunks(pdf, nchunks)
+        at = pyarrow.concat_tables([pyarrow.Table.from_pandas(c) for c in chunks])
         mdf = from_arrow(at)
+        assert (
+            len(
+                mdf._query_compiler._modin_frame._partitions[0][0]
+                .get()
+                .column(0)
+                .chunks
+            )
+            == nchunks
+        )
         df_equals(mdf, pdf)
 
+        mdt = mdf.dtypes[0]
+        pdt = pdf.dtypes[0]
+        assert mdt == "category"
+        assert isinstance(mdt, pandas.CategoricalDtype)
+        assert pandas.api.types.is_categorical_dtype(mdt)
+        assert str(mdt) == str(pdt)
+
+        if type(mdt) != pandas.CategoricalDtype:
+            # This is a lazy proxy.
+            # Make sure the table is not materialized yet.
+            assert mdt._table is not None
path: modin/experimental/core/execution/native/implementations/hdk_on_native/test/test_dataframe.py
text: @user1:
why do we need this if-conditioning if we intend to test a lazy proxy here?
```suggestion
        assert type(mdt) == LazyProxyCategoricalDtype
        # This is a lazy proxy.
        # Make sure the table is not materialized yet.
        assert mdt._table is not None
```

@author:
> why do we need this if-conditioning if we intend to test a lazy proxy here?

I've added this condition here because, if the implementation is changed and this LazyProxyCategoricalDtype is not used any more, the test will fail.

diff_hunk: @@ -2556,3 +2566,58 @@ def _is_trivial_index(cls, index):
             and index.min() == 0
             and index.max() == len(index) - 1
         )
+
+
+class LazyProxyCategoricalDtype(pandas.CategoricalDtype):
+    """
+    Proxy class for lazily retrieving categorical dtypes from arrow tables.
+
+    Parameters
+    ----------
+    table : pyarrow.Table
+        Source table.
+    name : str
+        Column name.
+    """
+
+    def __init__(self, table: pyarrow.Table, name: str):
+        self._table = table
+        self._name = name
+        self._ordered = False
+        self._lazy_categories = None
+
+    def _new(self, table: pyarrow.Table, name: str) -> pandas.CategoricalDtype:
+        """
+        Create a new proxy, if either table or column name are different.
+
+        Parameters
+        ----------
+        table : pyarrow.Table
+            Source table.
+        name : str
+            Column name.
+
+        Returns
+        -------
+        pandas.CategoricalDtype
+        """
+        if self._table is None:
+            # The table has been materialized, we don't need a proxy anymore.
+            return pandas.CategoricalDtype(self.categories)
+        elif table is self._table and name == self._name:
+            return self
+        else:
+            return LazyProxyCategoricalDtype(table, name)
+
+    @property
+    def _categories(self):  # noqa: GL08
+        if self._table is not None:
+            chunks = self._table.column(self._name).chunks
+            cat = pandas.concat([chunk.dictionary.to_pandas() for chunk in chunks])
+            self._lazy_categories = self.validate_categories(cat.unique())
+            self._table = None  # The table is not required any more
+        return self._lazy_categories
+
+    @_categories.setter
+    def _set_categories(self, categories):  # noqa: GL08
+        self._lazy_categories = categories
path: modin/experimental/core/execution/native/implementations/hdk_on_native/dataframe/dataframe.py
text: @user1:
Don't we want to drop the pyarrow table right away once actual values were specified?

@author:
It's dropped on the line 2618, but you are right, I'll add it here either.

diff_hunk: @@ -2556,3 +2566,58 @@ def _is_trivial_index(cls, index):
             and index.min() == 0
             and index.max() == len(index) - 1
         )
+
+
+class LazyProxyCategoricalDtype(pandas.CategoricalDtype):
+    """
+    Proxy class for lazily retrieving categorical dtypes from arrow tables.
+
+    Parameters
+    ----------
+    table : pyarrow.Table
+        Source table.
+    name : str
+        Column name.
+    """
+
+    def __init__(self, table: pyarrow.Table, name: str):
+        self._table = table
+        self._name = name
+        self._ordered = False
+        self._lazy_categories = None
+
+    def _new(self, table: pyarrow.Table, name: str) -> pandas.CategoricalDtype:
+        """
+        Create a new proxy, if either table or column name are different.
+
+        Parameters
+        ----------
+        table : pyarrow.Table
+            Source table.
+        name : str
+            Column name.
+
+        Returns
+        -------
+        pandas.CategoricalDtype
+        """
+        if self._table is None:
+            # The table has been materialized, we don't need a proxy anymore.
+            return pandas.CategoricalDtype(self.categories)
+        elif table is self._table and name == self._name:
+            return self
+        else:
+            return LazyProxyCategoricalDtype(table, name)
path: modin/experimental/core/execution/native/implementations/hdk_on_native/dataframe/dataframe.py
text: @user1:
we use `None` as a special value in this class indicating that the lazy categories are not lazy anymore, let's make sure that this value won't be set by accident
```suggestion
        else:
            ErrorMessage.catch_bugs_and_request_email(
                failure_condition=table is None,
                extra_log="attempted to bind 'None' pyarrow table to a lazy category",
            )
            return LazyProxyCategoricalDtype(table, name)
```
p.s. you have to import `ErrorMessage` from [`modin.error_message`](https://github.com/modin-project/modin/blob/17e55993c952eb31fd453d98ead94a014f549115/modin/error_message.py#L20)

@author:
Thanks, I'll add this check to __init__()