Repository: modin-project/modin

Issue:
OmniSci backend with Arrow 3.0



Discussion:

Resolving patch:
diff --git a/modin/experimental/engines/omnisci_on_native/frame/omnisci_worker.py b/modin/experimental/engines/omnisci_on_native/frame/omnisci_worker.py
--- a/modin/experimental/engines/omnisci_on_native/frame/omnisci_worker.py
+++ b/modin/experimental/engines/omnisci_on_native/frame/omnisci_worker.py
@@ -132,40 +132,35 @@ def _genName(cls, name):
         # TODO: reword name in case of caller's mistake
         return name
 
-    @classmethod
-    def put_arrow_to_omnisci(cls, table, name=None):
+    @staticmethod
+    def cast_to_compatible_types(table):
         """
-        Import Arrow table to OmniSci backend.
+        Cast PyArrow table to be fully compatible with OmniSci.
 
         Parameters
         ----------
         table : pyarrow.Table
-            A table to import.
-        name : str, optional
-            A table name to use. None to generate a unique name.
+            Source table.
 
         Returns
         -------
-        str
-            Imported table name.
+        pyarrow.Table
+            Table with fully compatible types with OmniSci.
         """
-        name = cls._genName(name)
-
-        # Currently OmniSci doesn't support Arrow table import with
-        # dictionary columns. Here we cast dictionaries until support
-        # is in place.
-        # https://github.com/modin-project/modin/issues/1738
         schema = table.schema
         new_schema = schema
         need_cast = False
         new_cols = {}
         for i, field in enumerate(schema):
+            # Currently OmniSci doesn't support Arrow table import with
+            # dictionary columns. Here we cast dictionaries until support
+            # is in place.
+            # https://github.com/modin-project/modin/issues/1738
             if pa.types.is_dictionary(field.type):
                 # Conversion for dictionary of null type to string is not supported
                 # in Arrow. Build new column for this case for now.
                 if pa.types.is_null(field.type.value_type):
-                    mask_vals = np.full(table.num_rows, True, dtype=bool)
-                    mask = pa.array(mask_vals)
+                    mask = np.full(table.num_rows, True, dtype=bool)
                     new_col_data = np.empty(table.num_rows, dtype=str)
                     new_col = pa.array(new_col_data, pa.string(), mask)
                     new_cols[i] = new_col
@@ -175,6 +170,16 @@ def put_arrow_to_omnisci(cls, table, name=None):
                     field.name, pa.string(), field.nullable, field.metadata
                 )
                 new_schema = new_schema.set(i, new_field)
+            # OmniSci doesn't support importing Arrow's date type:
+            # https://github.com/omnisci/omniscidb/issues/678
+            elif pa.types.is_date(field.type):
+                # Arrow's date is the number of days since the UNIX-epoch, so we can convert it
+                # to a timestamp[s] (number of seconds since the UNIX-epoch) without losing precision
+                new_field = pa.field(
+                    field.name, pa.timestamp("s"), field.nullable, field.metadata
+                )
+                new_schema = new_schema.set(i, new_field)
+                need_cast = True
 
         for i, col in new_cols.items():
             table = table.set_column(i, new_schema[i], col)
@@ -182,6 +187,28 @@ def put_arrow_to_omnisci(cls, table, name=None):
         if need_cast:
             table = table.cast(new_schema)
 
+        return table
+
+    @classmethod
+    def put_arrow_to_omnisci(cls, table, name=None):
+        """
+        Import Arrow table to OmniSci engine.
+
+        Parameters
+        ----------
+        table : pyarrow.Table
+            A table to import.
+        name : str, optional
+            A table name to use. None to generate a unique name.
+
+        Returns
+        -------
+        str
+            Imported table name.
+        """
+        name = cls._genName(name)
+
+        table = cls.cast_to_compatible_types(table)
         fragment_size = OmnisciFragmentSize.get()
         if fragment_size is None:
             cpu_count = os.cpu_count()


Test patch:
diff --git a/modin/experimental/engines/omnisci_on_native/test/test_dataframe.py b/modin/experimental/engines/omnisci_on_native/test/test_dataframe.py
--- a/modin/experimental/engines/omnisci_on_native/test/test_dataframe.py
+++ b/modin/experimental/engines/omnisci_on_native/test/test_dataframe.py
@@ -38,6 +38,7 @@
     eval_general,
     df_equals_with_non_stable_indices,
 )
+from modin.utils import try_cast_to_pandas
 
 from modin.experimental.engines.omnisci_on_native.frame.partition_manager import (
     OmnisciOnNativeFramePartitionManager,
@@ -588,7 +589,17 @@ def concat(df1, df2, lib, **kwargs):
             df4 = df2.copy()
             return lib.concat([df1, df2, df3, df4])
 
-        run_and_compare(concat, data=self.data, data2=self.data2)
+        def sort_comparator(df1, df2):
+            """Sort and verify equality of the passed frames."""
+            # We sort values because order of rows in the 'union all' result is inconsistent in OmniSci
+            df1, df2 = (
+                try_cast_to_pandas(df).sort_values(df.columns[0]) for df in (df1, df2)
+            )
+            return df_equals(df1, df2)
+
+        run_and_compare(
+            concat, data=self.data, data2=self.data2, comparator=sort_comparator
+        )
 
     def test_concat_agg(self):
         def concat(lib, df1, df2):
diff --git a/modin/experimental/engines/omnisci_on_native/test/utils.py b/modin/experimental/engines/omnisci_on_native/test/utils.py
--- a/modin/experimental/engines/omnisci_on_native/test/utils.py
+++ b/modin/experimental/engines/omnisci_on_native/test/utils.py
@@ -13,7 +13,9 @@
 
 import pytest
 import modin.pandas as pd
+from modin.utils import try_cast_to_pandas
 import pandas
+from pandas.api.types import is_datetime64_any_dtype
 import pyarrow as pa
 
 from modin.pandas.test.utils import (
@@ -50,6 +52,9 @@ def eval_io(
     def omnisci_comparator(df1, df2):
         """Evaluate equality comparison of the passed frames after importing the Modin's one to OmniSci."""
         with ForceOmnisciImport(df1, df2):
+            # 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)
             comparator(df1, df2)
 
     general_eval_io(
@@ -66,6 +71,41 @@ def omnisci_comparator(df1, df2):
     )
 
 
+def align_datetime_dtypes(*dfs):
+    """
+    Make all of the passed frames have DateTime dtype for the same columns.
+
+    Cast column type of the certain frame to the DateTime type if any frame in
+    the `dfs` sequence has DateTime type for this column.
+
+    Parameters
+    ----------
+    *dfs : iterable of DataFrames
+        DataFrames to align DateTime dtypes.
+
+    Notes
+    -----
+    Passed Modin frames may be casted to pandas in the result.
+    """
+    datetime_cols = {}
+    for df in dfs:
+        for col, dtype in df.dtypes.items():
+            # If we already decided to cast this column to DateTime no more actions are needed
+            if col not in datetime_cols and is_datetime64_any_dtype(dtype):
+                datetime_cols[col] = dtype
+
+    casted_dfs = (
+        # OmniSci has difficulties with casting to certain dtypes (i.e. datetime64),
+        # so casting it to pandas before doing 'astype'
+        tuple(try_cast_to_pandas(df).astype(datetime_cols) for df in dfs)
+        # This is required so we don't try to cast empty OmniSci frames to pandas:
+        # https://github.com/modin-project/modin/issues/3428
+        if len(datetime_cols)
+        else dfs
+    )
+    return casted_dfs
+
+
 class ForceOmnisciImport:
     """
     Trigger import execution for Modin DataFrames obtained by OmniSci engine if already not.
diff --git a/modin/pandas/test/test_io.py b/modin/pandas/test/test_io.py
--- a/modin/pandas/test/test_io.py
+++ b/modin/pandas/test/test_io.py
@@ -46,7 +46,10 @@
 )
 
 if Backend.get() == "Omnisci":
-    from modin.experimental.engines.omnisci_on_native.test.utils import eval_io
+    from modin.experimental.engines.omnisci_on_native.test.utils import (
+        eval_io,
+        align_datetime_dtypes,
+    )
 else:
     from .utils import eval_io
 
@@ -1105,6 +1108,11 @@ def wrapped_read_csv(file, method):
             pytest.csvs_names["test_read_csv_regular"], method="modin"
         )
 
+        if Backend.get() == "Omnisci":
+            # Aligning DateTime dtypes because of the bug related to the `parse_dates` parameter:
+            # https://github.com/modin-project/modin/issues/3485
+            modin_df, pandas_df = align_datetime_dtypes(modin_df, pandas_df)
+
         df_equals(modin_df, pandas_df)
 
     @pytest.mark.parametrize(