Repository: modin-project/modin

Issue:
Implement usecols parameter for read_csv with OmniSci backend
**Is your feature request related to a problem? Please describe.**
A clear and concise description of what the problem is. What kind of performance improvements would you like to see with this new API?
According to pyarrow documentation `pyarrow.read_csv` supports `include_columns` (https://arrow.apache.org/docs/python/generated/pyarrow.csv.ConvertOptions.html#pyarrow.csv.ConvertOptions), it can be used for implementation of `usecols` parameter of `modin.read_csv` with OmniSci backend.



Discussion:

Resolving patch:
diff --git a/modin/experimental/engines/omnisci_on_ray/io.py b/modin/experimental/engines/omnisci_on_ray/io.py
--- a/modin/experimental/engines/omnisci_on_ray/io.py
+++ b/modin/experimental/engines/omnisci_on_ray/io.py
@@ -19,6 +19,9 @@
 from pyarrow.csv import read_csv, ParseOptions, ConvertOptions, ReadOptions
 import pyarrow as pa
 
+import pandas
+from pandas.io.parsers import _validate_usecols_arg
+
 
 class OmnisciOnRayIO(RayIO):
 
@@ -172,6 +175,8 @@ def read_csv(
                     "Specified a delimiter and delim_whitespace=True; you can only specify one."
                 )
 
+            usecols_md = cls._prepare_pyarrow_usecols(mykwargs)
+
             po = ParseOptions(
                 delimiter="\\s+" if delim_whitespace else delimiter,
                 quote_char=quotechar,
@@ -190,7 +195,7 @@ def read_csv(
                 # didn't passed explicitly as an array or a dict
                 timestamp_parsers=[""] if isinstance(parse_dates, bool) else None,
                 strings_can_be_null=None,
-                include_columns=None,
+                include_columns=usecols_md,
                 include_missing_columns=None,
                 auto_dict_encode=None,
                 auto_dict_max_cardinality=None,
@@ -229,3 +234,57 @@ def _dtype_to_arrow(cls, dtype):
             return pa.string()
         else:
             return pa.from_numpy_dtype(tname)
+
+    @classmethod
+    def _prepare_pyarrow_usecols(cls, read_csv_kwargs):
+        """
+        Define `usecols` parameter in the way pyarrow can process it.
+        ----------
+        read_csv_kwargs:
+                read_csv function parameters.
+
+        Returns
+        -------
+        usecols_md: list
+                Redefined `usecols` parameter.
+        """
+        usecols = read_csv_kwargs.get("usecols", None)
+        engine = read_csv_kwargs.get("engine", None)
+        usecols_md, usecols_names_dtypes = _validate_usecols_arg(usecols)
+        if usecols_md:
+            empty_pd_df = pandas.read_csv(
+                **dict(
+                    read_csv_kwargs,
+                    nrows=0,
+                    skipfooter=0,
+                    usecols=None,
+                    engine=None if engine == "arrow" else engine,
+                )
+            )
+            column_names = empty_pd_df.columns
+            if usecols_names_dtypes == "string":
+                if usecols_md.issubset(set(column_names)):
+                    # columns should be sorted because pandas doesn't preserve columns order
+                    usecols_md = [
+                        col_name for col_name in column_names if col_name in usecols_md
+                    ]
+                else:
+                    raise NotImplementedError(
+                        "values passed in the `usecols` parameter don't match columns names"
+                    )
+            elif usecols_names_dtypes == "integer":
+                # columns should be sorted because pandas doesn't preserve columns order
+                usecols_md = sorted(usecols_md)
+                if len(column_names) < usecols_md[-1]:
+                    raise NotImplementedError(
+                        "max usecols value is higher than the number of columns"
+                    )
+                usecols_md = [column_names[i] for i in usecols_md]
+            elif callable(usecols_md):
+                usecols_md = [
+                    col_name for col_name in column_names if usecols_md(col_name)
+                ]
+            else:
+                raise NotImplementedError("unsupported `usecols` parameter")
+
+        return usecols_md


Test patch:
diff --git a/modin/experimental/engines/omnisci_on_ray/test/test_dataframe.py b/modin/experimental/engines/omnisci_on_ray/test/test_dataframe.py
--- a/modin/experimental/engines/omnisci_on_ray/test/test_dataframe.py
+++ b/modin/experimental/engines/omnisci_on_ray/test/test_dataframe.py
@@ -328,6 +328,36 @@ def test_read_csv_datetime(
             parse_dates=parse_dates,
         )
 
+    @pytest.mark.parametrize("engine", [None, "arrow"])
+    @pytest.mark.parametrize(
+        "usecols",
+        [
+            None,
+            ["col1"],
+            ["col1", "col1"],
+            ["col1", "col2", "col6"],
+            ["col6", "col2", "col1"],
+            [0],
+            [0, 0],
+            [0, 1, 5],
+            [5, 1, 0],
+            lambda x: x in ["col1", "col2"],
+        ],
+    )
+    def test_read_csv_col_handling(
+        self,
+        engine,
+        usecols,
+    ):
+        eval_io(
+            fn_name="read_csv",
+            check_kwargs_callable=not callable(usecols),
+            md_extra_kwargs={"engine": engine},
+            # read_csv kwargs
+            filepath_or_buffer=pytest.csvs_names["test_read_csv_regular"],
+            usecols=usecols,
+        )
+
 
 class TestMasks:
     data = {