Repository: modin-project/modin

Issue:
PERF: HDK triggers LazyProxyCategoricalDtype materialization on merge
Before the merge, HDK checks dtypes and it triggers LazyProxyCategoricalDtype materialization.


Discussion:

Resolving patch:
diff --git a/modin/core/dataframe/pandas/metadata/dtypes.py b/modin/core/dataframe/pandas/metadata/dtypes.py
--- a/modin/core/dataframe/pandas/metadata/dtypes.py
+++ b/modin/core/dataframe/pandas/metadata/dtypes.py
@@ -12,6 +12,7 @@
 # governing permissions and limitations under the License.
 
 """Module contains class ``ModinDtypes``."""
+from typing import Union
 
 import pandas
 
@@ -202,7 +203,7 @@ def _update_proxy(self, parent, column_name):
             return self._build_proxy(parent, column_name, self._materializer)
 
     @classmethod
-    def _build_proxy(cls, parent, column_name, materializer):
+    def _build_proxy(cls, parent, column_name, materializer, dtype=None):
         """
         Construct a lazy proxy.
 
@@ -214,6 +215,8 @@ def _build_proxy(cls, parent, column_name, materializer):
             Column name of the categorical column in the source object.
         materializer : callable(parent, column_name) -> pandas.CategoricalDtype
             A function to call in order to extract categorical values.
+        dtype : dtype, optional
+            The categories dtype.
 
         Returns
         -------
@@ -223,8 +226,21 @@ def _build_proxy(cls, parent, column_name, materializer):
         result._parent = parent
         result._column_name = column_name
         result._materializer = materializer
+        result._dtype = dtype
         return result
 
+    def _get_dtype(self):
+        """
+        Get the categories dtype.
+
+        Returns
+        -------
+        dtype
+        """
+        if self._dtype is None:
+            self._dtype = self.categories.dtype
+        return self._dtype
+
     def __reduce__(self):
         """
         Serialize an object of this class.
@@ -265,6 +281,7 @@ def _categories(self, categories):
         self._categories_val = categories
         self._parent = None  # The parent is not required any more
         self._materializer = None
+        self._dtype = None
 
     @property
     def _is_materialized(self) -> bool:
@@ -286,3 +303,24 @@ def _materialize_categories(self):
         categoricals = self._materializer(self._parent, self._column_name)
         self._categories = categoricals.categories
         self._ordered = categoricals.ordered
+
+
+def get_categories_dtype(
+    cdt: Union[LazyProxyCategoricalDtype, pandas.CategoricalDtype]
+):
+    """
+    Get the categories dtype.
+
+    Parameters
+    ----------
+    cdt : LazyProxyCategoricalDtype or pandas.CategoricalDtype
+
+    Returns
+    -------
+    dtype
+    """
+    return (
+        cdt._get_dtype()
+        if isinstance(cdt, LazyProxyCategoricalDtype)
+        else cdt.categories.dtype
+    )
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
@@ -38,6 +38,7 @@
 )
 from modin.core.dataframe.pandas.dataframe.dataframe import PandasDataframe
 from modin.core.dataframe.pandas.metadata import LazyProxyCategoricalDtype
+from modin.core.dataframe.pandas.metadata.dtypes import get_categories_dtype
 from modin.core.dataframe.pandas.utils import concatenate
 from modin.error_message import ErrorMessage
 from modin.experimental.core.storage_formats.hdk.query_compiler import (
@@ -74,6 +75,7 @@
 from .utils import (
     ColNameCodec,
     arrow_to_pandas,
+    arrow_type_to_pandas,
     build_categorical_from_at,
     check_cols_to_join,
     check_join_supported,
@@ -1105,8 +1107,8 @@ def join(
             if isinstance(left_dt, pd.CategoricalDtype) and isinstance(
                 right_dt, pd.CategoricalDtype
             ):
-                left_dt = left_dt.categories.dtype
-                right_dt = right_dt.categories.dtype
+                left_dt = get_categories_dtype(left_dt)
+                right_dt = get_categories_dtype(right_dt)
             if not (
                 (is_integer_dtype(left_dt) and is_integer_dtype(right_dt))
                 or (is_string_dtype(left_dt) and is_string_dtype(right_dt))
@@ -2847,6 +2849,7 @@ def from_arrow(
                         parent=at,
                         column_name=col._name,
                         materializer=build_categorical_from_at,
+                        dtype=arrow_type_to_pandas(col.type.value_type),
                     )
                 )
             else:
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
--- 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
@@ -565,6 +565,23 @@ def mapper(at):
     return df
 
 
+def arrow_type_to_pandas(at: pa.lib.DataType):
+    """
+    Convert the specified arrow type to pandas dtype.
+
+    Parameters
+    ----------
+    at : pa.lib.DataType
+
+    Returns
+    -------
+    dtype
+    """
+    if at == pa.string():
+        return _get_dtype(str)
+    return at.to_pandas_dtype()
+
+
 class _CategoricalDtypeMapper:  # noqa: GL08
     @staticmethod
     def __from_arrow__(arr):  # noqa: GL08


Test patch: