Repository: modin-project/modin

Issue:
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>



Discussion:

Resolving patch:
diff --git a/modin/core/dataframe/pandas/dataframe/dataframe.py b/modin/core/dataframe/pandas/dataframe/dataframe.py
--- 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
--- 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
@@ -21,6 +21,7 @@
 from modin.experimental.core.storage_formats.hdk.query_compiler import (
     DFAlgQueryCompiler,
 )
+from .utils import LazyProxyCategoricalDtype
 from ..partitioning.partition_manager import HdkOnNativeDataframePartitionManager
 
 from pandas.core.indexes.api import ensure_index, Index, MultiIndex, RangeIndex
@@ -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,
         )
diff --git a/modin/experimental/core/execution/native/implementations/hdk_on_native/dataframe/utils.py b/modin/experimental/core/execution/native/implementations/hdk_on_native/dataframe/utils.py
new file mode 100644
--- /dev/null
+++ b/modin/experimental/core/execution/native/implementations/hdk_on_native/dataframe/utils.py
@@ -0,0 +1,80 @@
+# Licensed to Modin Development Team under one or more contributor license agreements.
+# See the NOTICE file distributed with this work for additional information regarding
+# copyright ownership.  The Modin Development Team licenses this file to you under the
+# Apache License, Version 2.0 (the "License"); you may not use this file except in
+# compliance with the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software distributed under
+# the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific language
+# governing permissions and limitations under the License.
+
+"""Utilities for internal use by the ``HdkOnNativeDataframe``."""
+
+
+import pandas
+import pyarrow
+
+from modin.error_message import ErrorMessage
+
+
+class LazyProxyCategoricalDtype(pandas.CategoricalDtype):
+    """
+    Proxy class for lazily retrieving categorical dtypes from arrow tables.
+
+    Parameters
+    ----------
+    table : pyarrow.Table
+        Source table.
+    column_name : str
+        Column name.
+    """
+
+    def __init__(self, table: pyarrow.Table, column_name: str):
+        ErrorMessage.catch_bugs_and_request_email(
+            failure_condition=table is None,
+            extra_log="attempted to bind 'None' pyarrow table to a lazy category",
+        )
+        self._table = table
+        self._column_name = column_name
+        self._ordered = False
+        self._lazy_categories = None
+
+    def _new(self, table: pyarrow.Table, column_name: str) -> pandas.CategoricalDtype:
+        """
+        Create a new proxy, if either table or column name are different.
+
+        Parameters
+        ----------
+        table : pyarrow.Table
+            Source table.
+        column_name : str
+            Column name.
+
+        Returns
+        -------
+        pandas.CategoricalDtype or LazyProxyCategoricalDtype
+        """
+        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 column_name == self._column_name:
+            return self
+        else:
+            return LazyProxyCategoricalDtype(table, column_name)
+
+    @property
+    def _categories(self):  # noqa: GL08
+        if self._table is not None:
+            chunks = self._table.column(self._column_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
+        self._table = None


Test patch:
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
--- 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,39 @@ 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)
+        at = mdf._query_compiler._modin_frame._partitions[0][0].get()
+        assert len(at.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)
+
+        # Make sure the lazy proxy dtype is not materialized yet.
+        assert type(mdt) != pandas.CategoricalDtype
+        assert mdt._table is not None
+        assert mdt._new(at, at.column(0)._name) is mdt
+        assert mdt._new(at, at.column(2)._name) is not mdt
+        assert type(mdt._new(at, at.column(2)._name)) != pandas.CategoricalDtype
+
+        assert mdt == pdt
+        assert pdt == mdt
+        assert repr(mdt) == repr(pdt)
+
+        # Should be materialized now
+        assert type(mdt._new(at, at.column(2)._name)) == pandas.CategoricalDtype
+
 
 if __name__ == "__main__":
     pytest.main(["-v", __file__])