dask · jrbourbeau · Apr 29, 2021 · Mar 12, 2021 · Mar 12, 2021 · Mar 12, 2021
diff --git a/dask/blockwise.py b/dask/blockwise.py
diff --git a/dask/dataframe/io/csv.py b/dask/dataframe/io/csv.py
@@ -1,7 +1,11 @@
+import copy
 from collections.abc import Mapping
 from io import BytesIO
 from warnings import catch_warnings, simplefilter, warn
 
+from ...highlevelgraph import HighLevelGraph
+from ...layers import DataFrameIOLayer
+
 try:
     import psutil
 except ImportError:
@@ -32,86 +36,82 @@
 from ..utils import clear_known_categories
 
 
-class CSVSubgraph(Mapping):
+class CSVFunctionWrapper:
     """
-    Subgraph for reading CSV files.
+    CSV Function-Wrapper Class
+    Reads CSV data from disk to produce a partition (given a key).
     """
 
     def __init__(
         self,
-        name,
-        reader,
-        blocks,
-        is_first,
+        columns,
+        colname,
         head,
         header,
-        kwargs,
+        reader,
         dtypes,
-        columns,
         enforce,
-        path,
+        kwargs,
     ):
-        self.name = name
+        self.full_columns = columns
+        self.colname = colname
+        self.head = head
+        self.header = header
         self.reader = reader
-        self.blocks = blocks
-        self.is_first = is_first
-        self.head = head  # example pandas DF for metadata
-        self.header = header  # prepend to all blocks
-        self.kwargs = kwargs
         self.dtypes = dtypes
-        self.columns = columns
         self.enforce = enforce
-        self.colname, self.paths = path or (None, None)
+        self.kwargs = kwargs
+        self.columns = None  # Used to pass `usecols`
 
-    def __getitem__(self, key):
-        try:
-            name, i = key
-        except ValueError:
-            # too many / few values to unpack
-            raise KeyError(key) from None
+    def project_columns(self, columns):
+        """Return a new CSVFunctionWrapper object with
+        a sub-column projection.
+        """
+        if columns == self.columns:
+            return self
+        func = copy.deepcopy(self)
+        func.columns = columns
+        return func
 
-        if name != self.name:
-            raise KeyError(key)
+    def __call__(self, part):
 
-        if i < 0 or i >= len(self.blocks):
-            raise KeyError(key)
+        # Part will be a 3-element tuple
+        block, path, is_first = part
 
-        block = self.blocks[i]
-        if self.paths is not None:
+        # Construct `path_info`
+        if path is not None:
             path_info = (
                 self.colname,
-                self.paths[i],
+                path,
                 sorted(list(self.head[self.colname].cat.categories)),
             )
         else:
             path_info = None
 
+        # Deal with arguments that are special
+        # for the first block of each file
         write_header = False
         rest_kwargs = self.kwargs.copy()
-        if not self.is_first[i]:
+        if self.columns is not None:
+            if rest_kwargs.get("usecols", None) is None:
+                rest_kwargs["usecols"] = self.columns
+        if not is_first:
             write_header = True
             rest_kwargs.pop("skiprows", None)
 
-        return (
-            pandas_read_text,
+        # Call `pandas_read_text`
+        return pandas_read_text(
             self.reader,
             block,
             self.header,
             rest_kwargs,
             self.dtypes,
-            self.columns,
+            self.full_columns,
             write_header,
             self.enforce,
             path_info,
         )
 
-    def __len__(self):
-        return len(self.blocks)
-
-    def __iter__(self):
-        for i in range(len(self)):
-            yield (self.name, i)
-
 
 def pandas_read_text(
     reader,
@@ -325,8 +325,6 @@ def text_blocks_to_pandas(
     # Create mask of first blocks from nested block_lists
     is_first = tuple(block_mask(block_lists))
 
-    name = "read-csv-" + tokenize(reader, columns, enforce, head, blocksize)
-
     if path:
         colname, path_converter = path
         paths = [b[1].path for b in blocks]
@@ -344,21 +342,40 @@ def text_blocks_to_pandas(
     if len(unknown_categoricals):
         head = clear_known_categories(head, cols=unknown_categoricals)
 
-    subgraph = CSVSubgraph(
+    # Define parts
+    parts = []
+    colname, paths = path or (None, None)
+    for i in range(len(blocks)):
+        parts.append(
+            [
+                blocks[i],
+                paths[i] if paths else None,
+                is_first[i],
+            ]
+        )
+
+    # Create Blockwise layer
+    label = "read-csv-"
+    name = label + tokenize(reader, columns, enforce, head, blocksize)
+    layer = DataFrameIOLayer(
         name,
-        reader,
-        blocks,
-        is_first,
-        head,
-        header,
-        kwargs,
-        dtypes,
         columns,
-        enforce,
-        path,
+        parts,
+        CSVFunctionWrapper(
+            columns,
+            colname,
+            head,
+            header,
+            reader,
+            dtypes,
+            enforce,
+            kwargs,
+        ),
+        label=label,
+        produces_tasks=True,
     )
-
-    return new_dd_object(subgraph, name, head, (None,) * (len(blocks) + 1))
+    graph = HighLevelGraph({name: layer}, {name: set()})
+    return new_dd_object(graph, name, head, (None,) * (len(blocks) + 1))
 
 
 def block_mask(block_lists):

diff --git a/dask/dataframe/io/orc.py b/dask/dataframe/io/orc.py
@@ -1,15 +1,49 @@
+import copy
 from distutils.version import LooseVersion
 
 from fsspec.core import get_fs_token_paths
 
 from ...base import tokenize
+from ...highlevelgraph import HighLevelGraph
+from ...layers import DataFrameIOLayer
 from ...utils import import_required
 from ..core import DataFrame
 from .utils import _get_pyarrow_dtypes, _meta_from_dtypes
 
 __all__ = ("read_orc",)
 
 
+class ORCFunctionWrapper:
+    """
+    ORC Function-Wrapper Class
+    Reads ORC data from disk to produce a partition.
+    """
+
+    def __init__(self, fs, columns, schema):
+        self.fs = fs
+        self.columns = columns
+        self.schema = schema
+
+    def project_columns(self, columns):
+        """Return a new ORCFunctionWrapper object with
+        a sub-column projection.
+        """
+        if columns == self.columns:
+            return self
+        func = copy.deepcopy(self)
+        func.columns = columns
+        return func
+
+    def __call__(self, stripe_info):
+        path, stripe = stripe_info
+        return _read_orc_stripe(
+            self.fs,
+            path,
+            stripe,
+            list(self.schema) if self.columns is None else self.columns,
+        )
+
+
 def _read_orc_stripe(fs, path, stripe, columns=None):
     """Pull out specific data from specific part of ORC file"""
     orc = import_required("pyarrow.orc", "Please install pyarrow >= 0.9.0")
@@ -26,7 +60,6 @@ def _read_orc_stripe(fs, path, stripe, columns=None):
 
 def read_orc(path, columns=None, storage_options=None):
     """Read dataframe from ORC file(s)
-
     Parameters
     ----------
     path: str or list(str)
@@ -36,11 +69,9 @@ def read_orc(path, columns=None, storage_options=None):
         Columns to load. If None, loads all.
     storage_options: None or dict
         Further parameters to pass to the bytes backend.
-
     Returns
     -------
     Dask.DataFrame (even if there is only one column)
-
     Examples
     --------
     >>> df = dd.read_orc('https://github.com/apache/orc/raw/'
@@ -64,32 +95,36 @@ def read_orc(path, columns=None, storage_options=None):
         path, mode="rb", storage_options=storage_options
     )
     schema = None
-    nstripes_per_file = []
+    parts = []
     for path in paths:
         with fs.open(path, "rb") as f:
             o = orc.ORCFile(f)
             if schema is None:
                 schema = o.schema
             elif schema != o.schema:
                 raise ValueError("Incompatible schemas while parsing ORC files")
-            nstripes_per_file.append(o.nstripes)
+        for stripe in range(o.nstripes):
+            parts.append((path, stripe))
     schema = _get_pyarrow_dtypes(schema, categories=None)
     if columns is not None:
         ex = set(columns) - set(schema)
         if ex:
             raise ValueError(
                 "Requested columns (%s) not in schema (%s)" % (ex, set(schema))
             )
-    else:
-        columns = list(schema)
-    meta = _meta_from_dtypes(columns, schema, [], [])
 
-    name = "read-orc-" + tokenize(fs_token, path, columns)
-    dsk = {}
-    N = 0
-    for path, n in zip(paths, nstripes_per_file):
-        for stripe in range(n):
-            dsk[(name, N)] = (_read_orc_stripe, fs, path, stripe, columns)
-            N += 1
+    # Create Blockwise layer
+    label = "read-orc-"
+    output_name = label + tokenize(fs_token, path, columns)
+    layer = DataFrameIOLayer(
+        output_name,
+        columns,
+        parts,
+        ORCFunctionWrapper(fs, columns, schema),
+        label=label,
+    )
 
-    return DataFrame(dsk, name, meta, [None] * (len(dsk) + 1))
+    columns = list(schema) if columns is None else columns
+    meta = _meta_from_dtypes(columns, schema, [], [])
+    graph = HighLevelGraph({output_name: layer}, {output_name: set()})
+    return DataFrame(graph, output_name, meta, [None] * (len(parts) + 1))
diff --git a/dask/dataframe/io/parquet/arrow.py b/dask/dataframe/io/parquet/arrow.py
@@ -1472,11 +1472,13 @@ def _process_metadata(
             filters,
         )
 
-        # Check if we need to pass a fragment for each output partition
-        read_from_paths = read_from_paths or False
+        # Check if we need to pass a fragment for each output partition.
+        # By default, we will avoid passing fragments in the graph unless
+        # the user has specified `read_from_paths=False`
+        partitions = partition_info.get("partitions", None)
         pass_frags = (
             filters
-            and (not read_from_paths)
+            and (read_from_paths is False)
             and _need_fragments(filters, partition_info.get("partition_keys", None))
         )
 
@@ -1492,7 +1494,7 @@ def _process_metadata(
             make_part_kwargs={
                 "fs": fs,
                 "partition_keys": partition_info.get("partition_keys", None),
-                "partition_obj": partition_info.get("partitions", None),
+                "partition_obj": partitions,
                 "data_path": data_path,
                 "frag_map": frag_map if pass_frags else None,
             },
@@ -1501,7 +1503,7 @@ def _process_metadata(
         # Add common kwargs
         common_kwargs = {
             "partitioning": partition_info["partitioning"],
-            "partitions": partition_info["partitions"],
+            "partitions": partitions,
             "categories": categories,
             "filters": filters,
             "schema": schema,