Repository: NVIDIA-Merlin/NVTabular

Issue:
[FEA] Implement CSV engine for `DaskDataset`
<a href="https://github.com/rjzamora"><img src="https://avatars0.githubusercontent.com/u/20461013?v=4" align="left" width="96" height="96" hspace="10"></img></a> **Issue by [rjzamora](https://github.com/rjzamora)**
_Thursday May 28, 2020 at 17:00 GMT_
_Originally opened as https://github.com/rapidsai/recsys/issues/181_

----

**Is your feature request related to a problem? Please describe.**
#136 introduces a new `DaskDataset` API for managing IO needs within a `DaskWorkflow`.  This API only supports `"parquet"` at the moment.  A similar engine for "csv" needs to be implemented.

**Describe the solution you'd like**
We need a `CSVDatasetEngine` implementation in `nvtabular/dask/io.py`.



Discussion:

Resolving patch:
diff --git a/nvtabular/dask/io.py b/nvtabular/dask/io.py
--- a/nvtabular/dask/io.py
+++ b/nvtabular/dask/io.py
@@ -20,6 +20,7 @@
 
 import cudf
 import cupy
+import dask_cudf
 import numba.cuda as cuda
 import numpy as np
 import pyarrow.parquet as pq
@@ -145,7 +146,9 @@ class DaskDataset:
     -----------
     path : str or list of str
         Dataset path (or list of paths). If string, should specify
-        a specific file or directory path.
+        a specific file or directory path. If this is a directory
+        path, the directory structure must be flat (nested directories
+        are not yet supported).
     engine : str or DatasetEngine
         DatasetEngine object or string identifier of engine. Current
         string options include: ("parquet").
@@ -159,17 +162,28 @@ class DaskDataset:
         to GPU memory capacity). Ignored if part_size is passed
         directly. Note that the underlying engine may allow other
         custom kwargs to override this argument.
+    storage_options: None or dict
+        Further parameters to pass to the bytes backend.
     **kwargs :
         Other arguments to be passed to DatasetEngine.
     """
 
-    def __init__(self, path, engine, part_size=None, part_mem_fraction=0.125, **kwargs):
+    def __init__(
+        self,
+        path,
+        engine=None,
+        part_size=None,
+        part_mem_fraction=None,
+        storage_options=None,
+        **kwargs,
+    ):
 
         if part_size:
             # If a specific partition size is given, use it directly
             part_size = parse_bytes(part_size)
         else:
             # If a fractional partition size is given, calculate part_size
+            part_mem_fraction = part_mem_fraction or 0.125
             assert part_mem_fraction > 0.0 and part_mem_fraction < 1.0
             if part_mem_fraction > 0.25:
                 warnings.warn(
@@ -178,43 +192,57 @@ def __init__(self, path, engine, part_size=None, part_mem_fraction=0.125, **kwar
                 )
             part_size = int(cuda.current_context().get_memory_info()[1] * part_mem_fraction)
 
+        # Engine-agnostic path handling
+        if hasattr(path, "name"):
+            path = stringify_path(path)
+        storage_options = storage_options or {}
+        fs, fs_token, paths = get_fs_token_paths(path, mode="rb", storage_options=storage_options)
+        paths = sorted(paths, key=natural_sort_key)
+
+        # If engine is not provided, try to infer from end of paths[0]
+        if engine is None:
+            engine = paths[0].split(".")[-1]
+
         if isinstance(engine, str):
             if engine == "parquet":
-                self.engine = ParquetDatasetEngine(path, part_size, **kwargs)
+                self.engine = ParquetDatasetEngine(paths, part_size, fs, fs_token, **kwargs)
+            elif engine == "csv":
+                self.engine = CSVDatasetEngine(paths, part_size, fs, fs_token, **kwargs)
             else:
-                raise ValueError("Only parquet supported for now")
+                raise ValueError("Only parquet and csv supported (for now).")
         else:
-            self.engine = engine(path, part_size, **kwargs)
-
-    def meta_empty(self, columns=None):
-        return self.engine.meta_empty(columns=columns)
+            self.engine = engine(paths, part_size, fs, fs_token, **kwargs)
 
     def to_ddf(self, columns=None):
         return self.engine.to_ddf(columns=columns)
 
 
 class DatasetEngine:
-    """ DaskDataset Class
-        Converts dataset `pieces` to a dask_cudf DataFrame
+    """ DatasetEngine Class
+
+        Base class for Dask-powered IO engines. Engines must provide
+        a ``to_ddf`` method.
     """
 
-    def __init__(self, path, part_size):
-        if hasattr(path, "name"):
-            path = stringify_path(path)
-        fs, _, paths = get_fs_token_paths(path, mode="rb")
-        self.fs = fs
-        self.paths = sorted(paths, key=natural_sort_key)
+    def __init__(self, paths, part_size, fs, fs_token):
+        self.paths = paths
         self.part_size = part_size
-
-    def meta_empty(self, columns=None):
-        raise NotImplementedError(""" Return an empty cudf.DataFrame with the correct schema """)
+        self.fs = fs
+        self.fs_token = fs_token
 
     def to_ddf(self, columns=None):
         raise NotImplementedError(""" Return a dask_cudf.DataFrame """)
 
 
 class ParquetDatasetEngine(DatasetEngine):
+    """ ParquetDatasetEngine
+
+        Dask-based version of cudf.read_parquet.
+    """
+
     def __init__(self, *args, row_groups_per_part=None):
+        # TODO: Improve dask_cudf.read_parquet performance so that
+        # this class can be slimmed down.
         super().__init__(*args)
         self._metadata, self._base = self.get_metadata()
         self._pieces = None
@@ -308,7 +336,7 @@ def meta_empty(self, columns=None):
 
     def to_ddf(self, columns=None):
         pieces = self.pieces
-        name = "parquet-to-ddf-" + tokenize(pieces, columns)
+        name = "parquet-to-ddf-" + tokenize(self.fs_token, pieces, columns)
         dsk = {
             (name, p): (ParquetDatasetEngine.read_piece, piece, columns)
             for p, piece in enumerate(pieces)
@@ -316,3 +344,25 @@ def to_ddf(self, columns=None):
         meta = self.meta_empty(columns=columns)
         divisions = [None] * (len(pieces) + 1)
         return new_dd_object(dsk, name, meta, divisions)
+
+
+class CSVDatasetEngine(DatasetEngine):
+    """ CSVDatasetEngine
+
+        Thin wrapper around dask_cudf.read_csv.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args)
+        self._meta = {}
+        self.names = kwargs.pop("names", None)
+        self.csv_kwargs = kwargs
+        # CSV reader needs a list of files
+        # (Assume flat directory structure if this is a dir)
+        if len(self.paths) == 1 and self.fs.isdir(self.paths[0]):
+            self.paths = self.fs.glob(self.fs.sep.join([self.paths[0], "*"]))
+
+    def to_ddf(self, columns=None):
+        return dask_cudf.read_csv(
+            self.paths, names=self.names, chunksize=self.part_size, **self.csv_kwargs
+        )[columns]


Test patch:
diff --git a/tests/unit/test_dask_nvt.py b/tests/unit/test_dask_nvt.py
--- a/tests/unit/test_dask_nvt.py
+++ b/tests/unit/test_dask_nvt.py
@@ -26,7 +26,7 @@
 
 import nvtabular.ops as ops
 from nvtabular import DaskDataset, Workflow
-from tests.conftest import mycols_pq
+from tests.conftest import allcols_csv, mycols_csv, mycols_pq
 
 # LocalCluster Client Fixture
 client = None
@@ -57,16 +57,29 @@ def _dummy_op_logic(gdf, target_columns, _id="dummy", **kwargs):
     return new_gdf
 
 
-@pytest.mark.parametrize("engine", ["parquet"])
+@pytest.mark.parametrize("part_mem_fraction", [0.01, None])
+@pytest.mark.parametrize("engine", ["parquet", "csv", "csv-no-header"])
 @pytest.mark.parametrize("freq_threshold", [0, 5])
-def test_dask_workflow_api_dlrm(dask_cluster, tmpdir, datasets, freq_threshold, engine):
+def test_dask_workflow_api_dlrm(
+    dask_cluster, tmpdir, datasets, freq_threshold, part_mem_fraction, engine
+):
 
     paths = glob.glob(str(datasets[engine]) + "/*." + engine.split("-")[0])
-    df1 = cudf.read_parquet(paths[0])[mycols_pq]
-    df2 = cudf.read_parquet(paths[1])[mycols_pq]
+    if engine == "parquet":
+        df1 = cudf.read_parquet(paths[0])[mycols_pq]
+        df2 = cudf.read_parquet(paths[1])[mycols_pq]
+    elif engine == "csv":
+        df1 = cudf.read_csv(paths[0], header=0)[mycols_csv]
+        df2 = cudf.read_csv(paths[1], header=0)[mycols_csv]
+    else:
+        df1 = cudf.read_csv(paths[0], names=allcols_csv)[mycols_csv]
+        df2 = cudf.read_csv(paths[1], names=allcols_csv)[mycols_csv]
     df0 = cudf.concat([df1, df2], axis=0)
 
-    cat_names = ["name-cat", "name-string"]
+    if engine == "parquet":
+        cat_names = ["name-cat", "name-string"]
+    else:
+        cat_names = ["name-string"]
     cont_names = ["x", "y", "id"]
     label_name = ["label"]
 
@@ -78,7 +91,10 @@ def test_dask_workflow_api_dlrm(dask_cluster, tmpdir, datasets, freq_threshold,
     processor.add_preprocess(ops.Categorify(freq_threshold=freq_threshold, out_path=str(tmpdir)))
     processor.finalize()
 
-    dataset = DaskDataset(paths, engine)
+    if engine in ("parquet", "csv"):
+        dataset = DaskDataset(paths, part_mem_fraction=part_mem_fraction)
+    else:
+        dataset = DaskDataset(paths, names=allcols_csv, part_mem_fraction=part_mem_fraction)
     processor.apply(dataset, output_path=str(tmpdir))
     result = processor.get_ddf().compute()
 
@@ -90,12 +106,16 @@ def test_dask_workflow_api_dlrm(dask_cluster, tmpdir, datasets, freq_threshold,
 
     # Check category counts
     if freq_threshold == 0:
-        assert len(df0["name-cat"].unique()) == len(result["name-cat"].unique())
+        if engine == "parquet":
+            assert len(df0["name-cat"].unique()) == len(result["name-cat"].unique())
         assert len(df0["name-string"].unique()) == len(result["name-string"].unique())
 
         df0["_count"] = cupy.ones(len(df0))
         result["_count"] = cupy.ones(len(result))
-        for col in ["name-cat", "name-string"]:
+        cat_list = ["name-string"]
+        if engine == "parquet":
+            cat_list = ["name-cat", "name-string"]
+        for col in cat_list:
             expect = df0.groupby(col, dropna=False).count()["_count"].sort_values("_count")
             got = result.groupby(col, dropna=False).count()["_count"].sort_values("_count")
             assert_eq(expect, got, check_index=False)