Repository: modin-project/modin

Issue:
HDK: Do not keep reference to arrow table imported to HDK 
When the table is imported, it's not required any more and could be garbage collected.


Discussion:

Resolving patch:
diff --git a/docs/conf.py b/docs/conf.py
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -40,6 +40,8 @@ def noop_decorator(*args, **kwargs):
     sys.modules["pyhdk"].PyDbEngine = type("PyDbEngine", (object,), {})
 if not hasattr(sys.modules["pyhdk.hdk"], "HDK"):
     sys.modules["pyhdk.hdk"].HDK = type("HDK", (object,), {})
+if not hasattr(sys.modules["pyhdk.hdk"], "QueryNode"):
+    sys.modules["pyhdk.hdk"].QueryNode = type("QueryNode", (object,), {})
 if not hasattr(sys.modules["pyhdk.hdk"], "ExecutionResult"):
     sys.modules["pyhdk.hdk"].ExecutionResult = type("ExecutionResult", (object,), {})
 if not hasattr(sys.modules["pyhdk.hdk"], "RelAlgExecutor"):
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
@@ -53,7 +53,6 @@
     build_categorical_from_at,
 )
 from ..db_worker import DbTable
-from ..hdk_worker import ImportedTable
 from ..partitioning.partition_manager import HdkOnNativeDataframePartitionManager
 from modin.core.dataframe.pandas.metadata import LazyProxyCategoricalDtype
 from modin.error_message import ErrorMessage
@@ -511,8 +510,6 @@ def _maybe_update_proxies(self, dtypes, new_parent=None):
         if self._partitions is None:
             return
         table = self._partitions[0][0].get()
-        if isinstance(table, ImportedTable):
-            table = table.to_arrow()
         if isinstance(table, pyarrow.Table):
             super()._maybe_update_proxies(dtypes, new_parent=table)
 
diff --git a/modin/experimental/core/execution/native/implementations/hdk_on_native/hdk_worker.py b/modin/experimental/core/execution/native/implementations/hdk_on_native/hdk_worker.py
--- a/modin/experimental/core/execution/native/implementations/hdk_on_native/hdk_worker.py
+++ b/modin/experimental/core/execution/native/implementations/hdk_on_native/hdk_worker.py
@@ -12,12 +12,12 @@
 # governing permissions and limitations under the License.
 
 """Module provides ``HdkWorker`` class."""
-from typing import Optional, Tuple, List
+from typing import Optional, Tuple, List, Union
 
 import pyarrow as pa
 import os
 
-from pyhdk.hdk import HDK, ExecutionResult, RelAlgExecutor
+from pyhdk.hdk import HDK, QueryNode, ExecutionResult, RelAlgExecutor
 
 from .base_worker import DbTable, BaseDbWorker
 
@@ -25,18 +25,24 @@
 from modin.config import HdkLaunchParameters, OmnisciFragmentSize, HdkFragmentSize
 
 
-class ExecutionResultTable(DbTable):
+class HdkTable(DbTable):
     """
-    Represents an ExecutionResult table.
+    Represents a table in the HDK database.
 
     Parameters
     ----------
-    result : ExecutionResult
+    table : QueryNode or ExecutionResult
     """
 
-    def __init__(self, result: ExecutionResult):
-        self.name = result.table_name
-        self._result = result
+    def __init__(self, table: Union[QueryNode, ExecutionResult]):
+        self.name = table.table_name
+        self._table = table
+
+    def __del__(self):
+        """Drop table."""
+        # The ExecutionResults are cleared by HDK.
+        if not isinstance(self._table, ExecutionResult):
+            HdkWorker.dropTable(self.name)
 
     @property
     @_inherit_docstrings(DbTable.shape)
@@ -56,13 +62,11 @@ def column_names(self) -> List[str]:
 
     @_inherit_docstrings(DbTable.to_arrow)
     def to_arrow(self) -> pa.Table:
-        at = getattr(self, "_at", None)
-        if at is None:
-            at = self._result.to_arrow()
-            if (names := getattr(self, "_column_names", None)) is not None:
-                at = at.rename_columns(names)
-            self._at = at
-        return at
+        return (
+            self._table.to_arrow()
+            if isinstance(self._table, ExecutionResult)
+            else self._table.run().to_arrow()
+        )
 
     def scan(self):
         """
@@ -72,45 +76,14 @@ def scan(self):
         -------
         QueryNode
         """
+        if isinstance(self._table, QueryNode):
+            return self._table
         scan = getattr(self, "_scan", None)
         if scan is None:
-            self._scan = scan = HdkWorker._hdk().scan(self._result.table_name)
+            self._scan = scan = HdkWorker._hdk().scan(self.name)
         return scan
 
 
-class ImportedTable(DbTable):
-    """
-    Represents an imported arrow table.
-
-    Parameters
-    ----------
-    table : pa.Table
-    name : str
-    """
-
-    def __init__(self, table: pa.Table, name: str):
-        self.name = name
-        self._table = table
-
-    def __del__(self):
-        """Drop table."""
-        HdkWorker.dropTable(self.name)
-
-    @property
-    @_inherit_docstrings(DbTable.shape)
-    def shape(self) -> Tuple[int, int]:
-        return (self._table.num_rows, self._table.num_columns)
-
-    @property
-    @_inherit_docstrings(DbTable.column_names)
-    def column_names(self) -> List[str]:
-        return self._table.column_names
-
-    @_inherit_docstrings(DbTable.to_arrow)
-    def to_arrow(self) -> pa.Table:
-        return self._table
-
-
 @_inherit_docstrings(BaseDbWorker)
 class HdkWorker(BaseDbWorker):  # noqa: PR01
     """PyHDK based wrapper class for HDK storage format."""
@@ -138,17 +111,14 @@ def executeRA(cls, query: str, exec_calcite=False):
         else:
             ra = query
         ra_executor = RelAlgExecutor(hdk._executor, hdk._schema_mgr, hdk._data_mgr, ra)
-        return ExecutionResultTable(
-            ra_executor.execute(device_type=cls._preferred_device)
-        )
+        return HdkTable(ra_executor.execute(device_type=cls._preferred_device))
 
     @classmethod
     def import_arrow_table(cls, table: pa.Table, name: Optional[str] = None):
         name = cls._genName(name)
-        compat_table = cls.cast_to_compatible_types(table)
-        fragment_size = cls.compute_fragment_size(compat_table)
-        cls._hdk().import_arrow(compat_table, name, fragment_size)
-        return ImportedTable(table, name)
+        table = cls.cast_to_compatible_types(table)
+        fragment_size = cls.compute_fragment_size(table)
+        return HdkTable(cls._hdk().import_arrow(table, name, fragment_size))
 
     @classmethod
     def compute_fragment_size(cls, table):


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
@@ -224,9 +224,6 @@ def test_null_col(self, null_dtype):
             with ForceHdkImport(exp):
                 exp = to_pandas(exp)
             exp["c"] = exp["c"].astype("string")
-            # The arrow table contains empty strings, when reading as category.
-            assert all(v == "" for v in exp["c"])
-            exp["c"] = None
 
         df_equals(ref, exp)
 
diff --git a/modin/experimental/core/execution/native/implementations/hdk_on_native/test/utils.py b/modin/experimental/core/execution/native/implementations/hdk_on_native/test/utils.py
--- a/modin/experimental/core/execution/native/implementations/hdk_on_native/test/utils.py
+++ b/modin/experimental/core/execution/native/implementations/hdk_on_native/test/utils.py
@@ -52,6 +52,15 @@ def hdk_comparator(df1, df2, **kwargs):
             # Aligning DateTime dtypes because of the bug related to the `parse_dates` parameter:
             # https://github.com/modin-project/modin/issues/3485
             df1, df2 = align_datetime_dtypes(df1, df2)
+
+            # 1. Replace NA with empty strings. HDK treats empty strings and NA equally.
+            # 2. HdkWorker.cast_to_compatible_types() converts all categorical columns to string.
+            for dtype in ("object", "category"):
+                for df in (df1, df2):
+                    sdf = df.select_dtypes(dtype)
+                    if len(sdf.columns) != 0:
+                        sdf = sdf.fillna("") if dtype == "object" else sdf.astype(str)
+                        df[sdf.columns] = sdf[sdf.columns]
             comparator(df1, df2, **kwargs)
 
     general_eval_io(