Repository: NVIDIA-Merlin/NVTabular

Issue:
[FEA] Support reading datasets from Avro files

We should add support for reading from Avro datasets.



Discussion:
There is not dask_cudf support to read Avro files. I looked into implementing it using cudf and to avoid double work I asked in rapids channel if there are plans to add support in near future. Keith Kraus told me to create a FEA Request in GItHub: https://github.com/rapidsai/cudf/issues/6496

If they add support we can use it in a similar way that what we are doing with CSV. If they don't (or we don't want to wait) we can implement it in a similar way that we are doing with Parquet.

Waiting to see what rapids does.


Resolving patch:
diff --git a/nvtabular/io/avro.py b/nvtabular/io/avro.py
new file mode 100644
--- /dev/null
+++ b/nvtabular/io/avro.py
@@ -0,0 +1,175 @@
+#
+# Copyright (c) 2020, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import warnings
+
+import cudf
+import uavro as ua
+from dask.base import tokenize
+from dask.dataframe.core import new_dd_object
+
+from .dataset_engine import DatasetEngine
+
+
+class AvroDatasetEngine(DatasetEngine):
+    """AvroDatasetEngine
+
+    Uses `uavro` to decompose dataset into groups of avro blocks.
+    Uses `cudf` to create new partitions.
+    """
+
+    def __init__(self, paths, part_size, storage_options=None, **kwargs):
+        super().__init__(paths, part_size, storage_options)
+        if kwargs != {}:
+            raise ValueError("Unexpected AvroDatasetEngine argument(s).")
+        self.blocksize = part_size
+
+        # Avro 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):
+
+        # Get list of pieces for each output
+        pieces, meta = self.process_metadata(columns=columns)
+
+        # TODO: Remove warning and avoid use of uavro in read_partition when
+        # cudf#6529 is fixed (https://github.com/rapidsai/cudf/issues/6529)
+        if len(pieces) > len(self.paths):
+            warnings.warn(
+                "Row-subset selection in cudf avro reader is currently broken. "
+                "Using uavro engine until cudf#6529 is addressed. "
+                "EXPECT POOR PERFORMANCE!! (compared to cuio-based reader)"
+            )
+
+        # Construct collection
+        token = tokenize(self.fs, self.paths, self.part_size, columns)
+        read_avro_name = "read-avro-partition-" + token
+        dsk = {
+            (read_avro_name, i): (AvroDatasetEngine.read_partition, self.fs, piece, columns)
+            for i, piece in enumerate(pieces)
+        }
+        return new_dd_object(dsk, read_avro_name, meta.iloc[:0], [None] * (len(pieces) + 1))
+
+    def process_metadata(self, columns=None):
+
+        with open(self.paths[0], "rb") as fo:
+            header = ua.core.read_header(fo)
+
+            # Use first block for metadata
+            num_rows = header["blocks"][0]["nrows"]
+            file_byte_count = header["blocks"][0]["size"]
+            meta = cudf.io.read_avro(self.paths[0], skiprows=0, num_rows=num_rows)
+
+            # Convert the desired in-memory GPU size to the expected
+            # on-disk storage size (blocksize)
+            df_byte_count = meta.memory_usage(deep=True).sum()
+            self.blocksize = int(float(file_byte_count) / df_byte_count * self.part_size)
+
+        # Break apart files at the "Avro block" granularity
+        pieces = []
+        for path in self.paths:
+            file_size = self.fs.du(path)
+            if file_size > self.blocksize:
+                part_count = 0
+                with open(path, "rb") as fo:
+                    header = ua.core.read_header(fo)
+                    ua.core.scan_blocks(fo, header, file_size)
+                    blocks = header["blocks"]
+
+                    file_row_offset, part_row_count = 0, 0
+                    file_block_offset, part_block_count = 0, 0
+                    file_byte_offset, part_byte_count = blocks[0]["offset"], 0
+
+                    for i, block in enumerate(blocks):
+                        part_row_count += block["nrows"]
+                        part_block_count += 1
+                        part_byte_count += block["size"]
+                        if part_byte_count >= self.blocksize:
+                            pieces.append(
+                                {
+                                    "path": path,
+                                    "rows": (file_row_offset, part_row_count),
+                                    "blocks": (file_block_offset, part_block_count),
+                                    "bytes": (file_byte_offset, part_byte_count),
+                                }
+                            )
+                            part_count += 1
+                            file_row_offset += part_row_count
+                            file_block_offset += part_block_count
+                            file_byte_offset += part_byte_count
+                            part_row_count = part_block_count = part_byte_count = 0
+
+                    if part_block_count:
+                        pieces.append(
+                            {
+                                "path": path,
+                                "rows": (file_row_offset, part_row_count),
+                                "blocks": (file_block_offset, part_block_count),
+                                "bytes": (file_byte_offset, part_byte_count),
+                            }
+                        )
+                        part_count += 1
+                if part_count == 1:
+                    # No need to specify a byte range since we
+                    # will need to read the entire file anyway.
+                    pieces[-1] = {"path": pieces[-1]["path"]}
+            else:
+                pieces.append({"path": path})
+
+        return pieces, meta
+
+    @classmethod
+    def read_partition(cls, fs, piece, columns):
+
+        path = piece["path"]
+        if "rows" in piece:
+
+            # See: (https://github.com/rapidsai/cudf/issues/6529)
+            # Using `uavro` library for now. This means we must covert
+            # data to pandas, and then to cudf (which is much slower
+            # than `cudf.read_avro`). TODO: Once `num_rows` is fixed,
+            # this can be changed to:
+            #
+            #   skiprows, num_rows = piece["rows"]
+            #   df = cudf.io.read_avro(
+            #       path, skiprows=skiprows, num_rows=num_rows
+            #   )
+
+            block_offset, part_blocks = piece["blocks"]
+            file_size = fs.du(piece["path"])
+            with fs.open(piece["path"], "rb") as fo:
+                header = ua.core.read_header(fo)
+                ua.core.scan_blocks(fo, header, file_size)
+                header["blocks"] = header["blocks"][block_offset : block_offset + part_blocks]
+
+                # Adjust the total row count
+                nrows = 0
+                for block in header["blocks"]:
+                    nrows += block["nrows"]
+                header["nrows"] = nrows
+
+                # Read in as pandas and convert to cudf (avoid block scan)
+                df = cudf.from_pandas(
+                    ua.core.filelike_to_dataframe(fo, file_size, header, scan=False)
+                )
+        else:
+            df = cudf.io.read_avro(path)
+
+        # Deal with column selection
+        if columns is None:
+            columns = list(df.columns)
+        return df[columns]
diff --git a/nvtabular/io/dataset.py b/nvtabular/io/dataset.py
--- a/nvtabular/io/dataset.py
+++ b/nvtabular/io/dataset.py
@@ -146,7 +146,7 @@ class Dataset:
         directories are not yet supported).
     engine : str or DatasetEngine
         DatasetEngine object or string identifier of engine. Current
-        string options include: ("parquet", "csv"). This argument
+        string options include: ("parquet", "csv", "avro"). This argument
         is ignored if path_or_source is a DataFrame type.
     part_size : str or int
         Desired size (in bytes) of each Dask partition.
@@ -227,6 +227,17 @@ def __init__(
                     self.engine = CSVDatasetEngine(
                         paths, part_size, storage_options=storage_options, **kwargs
                     )
+                elif engine == "avro":
+                    try:
+                        from .avro import AvroDatasetEngine
+                    except ImportError:
+                        raise RuntimeError(
+                            "Failed to import AvroDatasetEngine. Make sure uavro is installed."
+                        )
+
+                    self.engine = AvroDatasetEngine(
+                        paths, part_size, storage_options=storage_options, **kwargs
+                    )
                 else:
                     raise ValueError("Only parquet and csv supported (for now).")
             else:


Test patch:
diff --git a/tests/unit/test_io.py b/tests/unit/test_io.py
--- a/tests/unit/test_io.py
+++ b/tests/unit/test_io.py
@@ -22,8 +22,10 @@
 import dask
 import dask_cudf
 import numpy as np
+import pandas as pd
 import pytest
 from dask.dataframe import assert_eq
+from dask.dataframe.io.demo import names as name_list
 
 import nvtabular as nvt
 import nvtabular.io
@@ -296,3 +298,62 @@ def test_parquet_lists(tmpdir, freq_threshold):
     df_out = cudf.read_parquet(out_paths)
 
     assert df_out["Authors"].to_arrow().to_pylist() == [[1], [1, 4], [2, 3], [3]]
+
+
+@pytest.mark.parametrize("part_size", [None, "1KB"])
+@pytest.mark.parametrize("size", [100, 5000])
+@pytest.mark.parametrize("nfiles", [1, 2])
+def test_avro_basic(tmpdir, part_size, size, nfiles):
+
+    # Require uavro and fastavro library.
+    # Note that fastavro is only required to write
+    # avro files for testing, while uavro is actually
+    # used by AvroDatasetEngine.
+    fa = pytest.importorskip("fastavro")
+    pytest.importorskip("uavro")
+
+    # Define avro schema
+    schema = fa.parse_schema(
+        {
+            "name": "avro.example.User",
+            "type": "record",
+            "fields": [
+                {"name": "name", "type": "string"},
+                {"name": "age", "type": "int"},
+            ],
+        }
+    )
+
+    # Write avro dataset with two files.
+    # Collect block and record (row) count while writing.
+    nblocks = 0
+    nrecords = 0
+    paths = [os.path.join(str(tmpdir), f"test.{i}.avro") for i in range(nfiles)]
+    records = []
+    for path in paths:
+        names = np.random.choice(name_list, size)
+        ages = np.random.randint(18, 100, size)
+        data = [{"name": names[i], "age": ages[i]} for i in range(size)]
+        with open(path, "wb") as f:
+            fa.writer(f, schema, data)
+        with open(path, "rb") as fo:
+            avro_reader = fa.block_reader(fo)
+            for block in avro_reader:
+                nrecords += block.num_records
+                nblocks += 1
+                records += list(block)
+    if nfiles == 1:
+        paths = paths[0]
+
+    # Read back with dask.dataframe
+    df = nvt.Dataset(paths, part_size=part_size, engine="avro").to_ddf()
+
+    # Check basic length and partition count
+    if part_size == "1KB":
+        assert df.npartitions == nblocks
+    assert len(df) == nrecords
+
+    # Full comparison
+    expect = pd.DataFrame.from_records(records)
+    expect["age"] = expect["age"].astype("int32")
+    assert_eq(df.compute().reset_index(drop=True), expect)