Repository: modin-project/modin

Issue:
[OmniSci] Support column renames in arrow execution



Discussion:

Resolving patch:
diff --git a/modin/experimental/engines/omnisci_on_native/frame/data.py b/modin/experimental/engines/omnisci_on_native/frame/data.py
--- a/modin/experimental/engines/omnisci_on_native/frame/data.py
+++ b/modin/experimental/engines/omnisci_on_native/frame/data.py
@@ -1675,7 +1675,11 @@ def _can_execute_arrow(self):
                 self._op.row_indices is None and self._op.input[0]._can_execute_arrow()
             )
         elif isinstance(self._op, TransformNode):
-            return self._op.is_drop() and self._op.input[0]._can_execute_arrow()
+            return (
+                not self._uses_rowid
+                and self._op.is_simple_select()
+                and self._op.input[0]._can_execute_arrow()
+            )
         elif isinstance(self._op, UnionNode):
             return all(frame._can_execute_arrow() for frame in self._op.input)
         else:
@@ -1699,20 +1703,20 @@ def _execute_arrow(self):
         elif isinstance(self._op, MaskNode):
             return self._op.input[0]._arrow_row_slice(self._op.row_numeric_idx)
         elif isinstance(self._op, TransformNode):
-            return self._op.input[0]._arrow_col_slice(set(self._op.exprs.keys()))
+            return self._op.input[0]._arrow_select(self._op.exprs)
         elif isinstance(self._op, UnionNode):
             return self._arrow_concat(self._op.input)
         else:
             raise RuntimeError(f"Unexpected op ({type(self._op)}) in _execute_arrow")
 
-    def _arrow_col_slice(self, new_columns):
+    def _arrow_select(self, exprs):
         """
         Perform column selection on the frame using Arrow API.
 
         Parameters
         ----------
-        new_columns : list of str
-            Columns to select.
+        exprs : dict
+            Select expressions.
 
         Returns
         -------
@@ -1720,9 +1724,21 @@ def _arrow_col_slice(self, new_columns):
             The resulting table.
         """
         table = self._execute_arrow()
-        return table.drop(
-            [f"F_{col}" for col in self._table_cols if col not in new_columns]
-        )
+        schema = table.schema
+
+        new_fields = []
+        new_columns = []
+
+        for col, expr in exprs.items():
+            field = schema.field(f"F_{expr.column}")
+            if col != expr.column:
+                field = field.with_name(f"F_{col}")
+            new_fields.append(field)
+            new_columns.append(table.column(f"F_{expr.column}"))
+
+        new_schema = pyarrow.schema(new_fields)
+
+        return pyarrow.Table.from_arrays(new_columns, schema=new_schema)
 
     def _arrow_row_slice(self, row_numeric_idx):
         """
diff --git a/modin/experimental/engines/omnisci_on_native/frame/df_algebra.py b/modin/experimental/engines/omnisci_on_native/frame/df_algebra.py
--- a/modin/experimental/engines/omnisci_on_native/frame/df_algebra.py
+++ b/modin/experimental/engines/omnisci_on_native/frame/df_algebra.py
@@ -591,29 +591,18 @@ def copy(self):
         """
         return TransformNode(self.input[0], self.exprs, self.keep_index)
 
-    def is_drop(self):
+    def is_simple_select(self):
         """
-        Check if transform node is a simple drop of columns.
+        Check if transform node is a simple selection.
 
-        Zero or more columns may be dropped. Remaining columns should
-        preserve original columns order.
+        Simple selection can only use InputRefExpr expressions.
 
         Returns
         -------
         bool
-            True for drop transformation and False otherwise.
-        """
-        col_iter = iter(self.input[0]._table_cols)
-        try:
-            for col, expr in self.exprs.items():
-                if not isinstance(expr, InputRefExpr) or expr.column != col:
-                    return False
-                while next(col_iter) != col:
-                    pass
-        except StopIteration:
-            return False
-
-        return True
+            True for simple select and False otherwise.
+        """
+        return all(isinstance(expr, InputRefExpr) for expr in self.exprs.values())
 
     def _prints(self, prefix):
         """


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
@@ -2008,5 +2008,33 @@ def test_shape_hint_detection_from_arrow(self):
         assert df._query_compiler._shape_hint == "column"
 
 
+class TestArrowExecution:
+    data1 = {
+        "a": [1, 2, 3],
+        "b": [3, 4, 5],
+        "c": [6, 7, 8],
+    }
+    data2 = {
+        "a": [1, 2, 3],
+        "d": [3, 4, 5],
+        "e": [6, 7, 8],
+    }
+
+    def test_drop_rename_concat(self):
+        def drop_rename_concat(df1, df2, lib, **kwargs):
+            df1 = df1.rename(columns={"a": "new_a", "c": "new_b"})
+            df1 = df1.drop(columns="b")
+            df2 = df2.rename(columns={"a": "new_a", "d": "new_b"})
+            df2 = df2.drop(columns="e")
+            return lib.concat([df1, df2], ignore_index=True)
+
+        run_and_compare(
+            drop_rename_concat,
+            data=self.data1,
+            data2=self.data2,
+            force_arrow_execute=True,
+        )
+
+
 if __name__ == "__main__":
     pytest.main(["-v", __file__])