Repository: dask/dask

Issue:
Add support for `add_prefix and `add_suffix`
I tend to use these before doing DataFrame merges so that my column names (of the right merged dataframe) retain some semantic notion of where they came from. Switched to Dask and noticed these don't exist.

Is that a conscious decision to exclude or would you be welcome to its addition?

I tried a brief google search and tried filtering through the issues on this repo, but forgive me if I've missed an existing discussion.

* [`pandas.DataFrame.add_prefix`](https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.add_prefix.html)
* [`pandas.DataFrame.add_suffix`](https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.add_suffix.html)

Discussion:

Patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
index 50424098f..85cbe8145 100644
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -1679,6 +1679,20 @@ Dask Name: {name}, {task} tasks"""
                 result.divisions = (self.columns.min(), self.columns.max())
             return handle_out(out, result)
 
+    @derived_from(pd.DataFrame)
+    def add_prefix(self, prefix):
+        res = self.map_partitions(M.add_prefix, prefix)
+        if self.known_divisions and is_series_like(self):
+            res.divisions = tuple(prefix + str(division) for division in self.divisions)
+        return res
+
+    @derived_from(pd.DataFrame)
+    def add_suffix(self, suffix):
+        res = self.map_partitions(M.add_suffix, suffix)
+        if self.known_divisions and is_series_like(self):
+            res.divisions = tuple(str(division) + suffix for division in self.divisions)
+        return res
+
     @derived_from(pd.DataFrame)
     def abs(self):
         _raise_if_object_series(self, "abs")
diff --git a/dask/dataframe/io/csv.py b/dask/dataframe/io/csv.py
index 076fe2254..748ab4ff3 100644
--- a/dask/dataframe/io/csv.py
+++ b/dask/dataframe/io/csv.py
@@ -1,4 +1,3 @@
-import copy
 from collections.abc import Mapping
 from io import BytesIO
 from warnings import catch_warnings, simplefilter, warn
@@ -72,11 +71,17 @@ class CSVFunctionWrapper:
         columns = [c for c in self.head.columns if c in columns]
         if columns == self.columns:
             return self
-        func = copy.deepcopy(self)
-        func.head = self.head[columns]
-        func.dtypes = {c: self.dtypes[c] for c in columns}
-        func.columns = columns
-        return func
+        return CSVFunctionWrapper(
+            self.full_columns,
+            columns,
+            self.colname,
+            self.head[columns],
+            self.header,
+            self.reader,
+            {c: self.dtypes[c] for c in columns},
+            self.enforce,
+            self.kwargs,
+        )
 
     def __call__(self, part):
 
diff --git a/dask/dataframe/io/demo.py b/dask/dataframe/io/demo.py
index d10bf069f..02bcdd90d 100644
--- a/dask/dataframe/io/demo.py
+++ b/dask/dataframe/io/demo.py
@@ -1,5 +1,3 @@
-import copy
-
 import numpy as np
 import pandas as pd
 
@@ -72,11 +70,11 @@ class MakeTimeseriesPart:
     Makes a timeseries partition.
     """
 
-    def __init__(self, dtypes, freq, kwargs):
-        self.dtypes = dtypes
+    def __init__(self, dtypes, freq, kwargs, columns=None):
+        self.columns = columns or list(dtypes.keys())
+        self.dtypes = {c: dtypes[c] for c in self.columns}
         self.freq = freq
         self.kwargs = kwargs
-        self.columns = None
 
     def project_columns(self, columns):
         """Return a new MakeTimeseriesPart object with
@@ -84,10 +82,12 @@ class MakeTimeseriesPart:
         """
         if columns == self.columns:
             return self
-        func = copy.deepcopy(self)
-        func.columns = columns
-        func.dtypes = {c: self.dtypes[c] for c in columns}
-        return func
+        return MakeTimeseriesPart(
+            self.dtypes,
+            self.freq,
+            self.kwargs,
+            columns=columns,
+        )
 
     def __call__(self, part):
         divisions, state_data = part
diff --git a/dask/dataframe/io/hdf.py b/dask/dataframe/io/hdf.py
index 57ad5f563..7737eed36 100644
--- a/dask/dataframe/io/hdf.py
+++ b/dask/dataframe/io/hdf.py
@@ -11,6 +11,8 @@ from tlz import merge
 from ... import config, multiprocessing
 from ...base import compute_as_if_collection, get_scheduler, tokenize
 from ...delayed import Delayed, delayed
+from ...highlevelgraph import HighLevelGraph
+from ...layers import DataFrameIOLayer
 from ...utils import get_scheduler_lock
 from ..core import DataFrame, new_dd_object
 from .io import _link
@@ -268,142 +270,41 @@ and stopping index per file, or starting and stopping index of the global
 dataset."""
 
 
-def _read_single_hdf(
-    path,
-    key,
-    start=0,
-    stop=None,
-    columns=None,
-    chunksize=int(1e6),
-    sorted_index=False,
-    lock=None,
-    mode="a",
-):
-    """
-    Read a single hdf file into a dask.dataframe. Used for each file in
-    read_hdf.
+class HDFFunctionWrapper:
     """
+    HDF5 Function-Wrapper Class
 
-    def get_keys_stops_divisions(path, key, stop, sorted_index, chunksize):
-        """
-        Get the "keys" or group identifiers which match the given key, which
-        can contain wildcards. This uses the hdf file identified by the
-        given path. Also get the index of the last row of data for each matched
-        key.
-        """
-        with pd.HDFStore(path, mode=mode) as hdf:
-            import glob
+    Reads HDF5 data from disk to produce a partition (given a key).
+    """
 
-            if not glob.has_magic(key):
-                keys = [key]
-            else:
-                keys = [k for k in hdf.keys() if fnmatch(k, key)]
-                # https://github.com/dask/dask/issues/5934
-                # TODO: remove this part if/when pandas copes with all keys
-                keys.extend(
-                    n._v_pathname
-                    for n in hdf._handle.walk_nodes("/", classname="Table")
-                    if fnmatch(n._v_pathname, key)
-                    and n._v_name != "table"
-                    and n._v_pathname not in keys
-                )
-            stops = []
-            divisions = []
-            for k in keys:
-                storer = hdf.get_storer(k)
-                if storer.format_type != "table":
-                    raise TypeError(dont_use_fixed_error_message)
-                if stop is None:
-                    stops.append(storer.nrows)
-                elif stop > storer.nrows:
-                    raise ValueError(
-                        "Stop keyword exceeds dataset number "
-                        "of rows ({})".format(storer.nrows)
-                    )
-                else:
-                    stops.append(stop)
-                if sorted_index:
-                    division = [
-                        storer.read_column("index", start=start, stop=start + 1)[0]
-                        for start in range(0, storer.nrows, chunksize)
-                    ]
-                    division_end = storer.read_column(
-                        "index", start=storer.nrows - 1, stop=storer.nrows
-                    )[0]
-
-                    division.append(division_end)
-                    divisions.append(division)
-                else:
-                    divisions.append(None)
-
-        return keys, stops, divisions
-
-    def one_path_one_key(path, key, start, stop, columns, chunksize, division, lock):
+    def __init__(self, columns, dim, lock, common_kwargs):
+        self.columns = columns
+        self.lock = lock
+        self.common_kwargs = common_kwargs
+        self.dim = dim
+        if columns and dim > 1:
+            self.common_kwargs = merge(common_kwargs, {"columns": columns})
+
+    def project_columns(self, columns):
+        """Return a new HDFFunctionWrapper object with
+        a sub-column projection.
         """
-        Get the data frame corresponding to one path and one key (which should
-        not contain any wildcards).
-        """
-        empty = pd.read_hdf(path, key, mode=mode, stop=0)
-        if columns is not None:
-            empty = empty[columns]
-
-        token = tokenize(
-            (path, os.path.getmtime(path), key, start, stop, empty, chunksize, division)
-        )
-        name = "read-hdf-" + token
-        if empty.ndim == 1:
-            base = {"name": empty.name, "mode": mode}
-        else:
-            base = {"columns": empty.columns, "mode": mode}
-
-        if start >= stop:
-            raise ValueError(
-                "Start row number ({}) is above or equal to stop "
-                "row number ({})".format(start, stop)
-            )
-
-        def update(s):
-            new = base.copy()
-            new.update({"start": s, "stop": s + chunksize})
-            return new
-
-        dsk = dict(
-            ((name, i), (_pd_read_hdf, path, key, lock, update(s)))
-            for i, s in enumerate(range(start, stop, chunksize))
-        )
-
-        if division:
-            divisions = division
-        else:
-            divisions = [None] * (len(dsk) + 1)
-
-        return new_dd_object(dsk, name, empty, divisions)
-
-    keys, stops, divisions = get_keys_stops_divisions(
-        path, key, stop, sorted_index, chunksize
-    )
-    if (start != 0 or stop is not None) and len(keys) > 1:
-        raise NotImplementedError(read_hdf_error_msg)
-    from ..multi import concat
-
-    return concat(
-        [
-            one_path_one_key(path, k, start, s, columns, chunksize, d, lock)
-            for k, s, d in zip(keys, stops, divisions)
-        ]
-    )
+        if columns == self.columns:
+            return self
+        return HDFFunctionWrapper(columns, self.dim, self.lock, self.common_kwargs)
 
+    def __call__(self, part):
+        """Read from hdf5 file with a lock"""
 
-def _pd_read_hdf(path, key, lock, kwargs):
-    """Read from hdf5 file with a lock"""
-    if lock:
-        lock.acquire()
-    try:
-        result = pd.read_hdf(path, key, **kwargs)
-    finally:
-        if lock:
-            lock.release()
-    return result
+        path, key, kwargs = part
+        if self.lock:
+            self.lock.acquire()
+        try:
+            result = pd.read_hdf(path, key, **merge(self.common_kwargs, kwargs))
+        finally:
+            if self.lock:
+                self.lock.release()
+        return result
 
 
 def read_hdf(
@@ -506,25 +407,140 @@ def read_hdf(
             "When assuming pre-partitioned data, data must be "
             "read in its entirety using the same chunksizes"
         )
-    from ..multi import concat
-
-    return concat(
-        [
-            _read_single_hdf(
-                path,
-                key,
-                start=start,
-                stop=stop,
-                columns=columns,
-                chunksize=chunksize,
-                sorted_index=sorted_index,
-                lock=lock,
-                mode=mode,
-            )
-            for path in paths
-        ]
+
+    # Build metadata
+    with pd.HDFStore(paths[0], mode=mode) as hdf:
+        meta_key = _expand_key(key, hdf)[0]
+    meta = pd.read_hdf(paths[0], meta_key, mode=mode, stop=0)
+    if columns is not None:
+        meta = meta[columns]
+
+    # Common kwargs
+    if meta.ndim == 1:
+        common_kwargs = {"name": meta.name, "mode": mode}
+    else:
+        common_kwargs = {"mode": mode}
+
+    # Build parts
+    parts, divisions = _build_parts(
+        paths, key, start, stop, chunksize, sorted_index, mode
     )
 
+    # Construct Layer and Collection
+    label = "read-hdf-"
+    name = label + tokenize(paths, key, start, stop, sorted_index, chunksize, mode)
+    layer = DataFrameIOLayer(
+        name,
+        columns,
+        parts,
+        HDFFunctionWrapper(columns, meta.ndim, lock, common_kwargs),
+        label=label,
+    )
+    graph = HighLevelGraph({name: layer}, {name: set()})
+    return new_dd_object(graph, name, meta, divisions)
+
+
+def _build_parts(paths, key, start, stop, chunksize, sorted_index, mode):
+    """
+    Build the list of partition inputs and divisions for read_hdf
+    """
+    parts = []
+    global_divisions = []
+    for path in paths:
+
+        keys, stops, divisions = _get_keys_stops_divisions(
+            path, key, stop, sorted_index, chunksize, mode
+        )
+
+        for k, stop, division in zip(keys, stops, divisions):
+
+            if division and global_divisions:
+                global_divisions = global_divisions[:-1] + division
+            elif division:
+                global_divisions = division
+
+            parts.extend(_one_path_one_key(path, k, start, stop, chunksize))
+
+    return parts, global_divisions or [None] * (len(parts) + 1)
+
+
+def _one_path_one_key(path, key, start, stop, chunksize):
+    """
+    Get the DataFrame corresponding to one path and one key (which
+    should not contain any wildcards).
+    """
+
+    if start >= stop:
+        raise ValueError(
+            "Start row number ({}) is above or equal to stop "
+            "row number ({})".format(start, stop)
+        )
+
+    return [
+        (path, key, {"start": s, "stop": s + chunksize})
+        for i, s in enumerate(range(start, stop, chunksize))
+    ]
+
+
+def _expand_key(key, hdf):
+    import glob
+
+    if not glob.has_magic(key):
+        keys = [key]
+    else:
+        keys = [k for k in hdf.keys() if fnmatch(k, key)]
+        # https://github.com/dask/dask/issues/5934
+        # TODO: remove this part if/when pandas copes with all keys
+        keys.extend(
+            n._v_pathname
+            for n in hdf._handle.walk_nodes("/", classname="Table")
+            if fnmatch(n._v_pathname, key)
+            and n._v_name != "table"
+            and n._v_pathname not in keys
+        )
+    return keys
+
+
+def _get_keys_stops_divisions(path, key, stop, sorted_index, chunksize, mode):
+    """
+    Get the "keys" or group identifiers which match the given key, which
+    can contain wildcards (see _expand_path). This uses the hdf file
+    identified by the given path. Also get the index of the last row of
+    data for each matched key.
+    """
+    with pd.HDFStore(path, mode=mode) as hdf:
+        stops = []
+        divisions = []
+        keys = _expand_key(key, hdf)
+        for k in keys:
+            storer = hdf.get_storer(k)
+            if storer.format_type != "table":
+                raise TypeError(dont_use_fixed_error_message)
+            if stop is None:
+                stops.append(storer.nrows)
+            elif stop > storer.nrows:
+                raise ValueError(
+                    "Stop keyword exceeds dataset number "
+                    "of rows ({})".format(storer.nrows)
+                )
+            else:
+                stops.append(stop)
+            if sorted_index:
+                division = [
+                    storer.read_column("index", start=start, stop=start + 1)[0]
+                    for start in range(0, storer.nrows, chunksize)
+                ]
+                division_end = storer.read_column(
+                    "index", start=storer.nrows - 1, stop=storer.nrows
+                )[0]
+
+                division.append(division_end)
+                divisions.append(division)
+            else:
+                divisions.append(None)
+
+    return keys, stops, divisions
+
 
 from ..core import _Frame
 
diff --git a/dask/dataframe/io/orc.py b/dask/dataframe/io/orc.py
index f448a7391..407403108 100644
--- a/dask/dataframe/io/orc.py
+++ b/dask/dataframe/io/orc.py
@@ -1,4 +1,3 @@
-import copy
 from distutils.version import LooseVersion
 
 from fsspec.core import get_fs_token_paths
@@ -30,9 +29,7 @@ class ORCFunctionWrapper:
         """
         if columns == self.columns:
             return self
-        func = copy.deepcopy(self)
-        func.columns = columns
-        return func
+        return ORCFunctionWrapper(self.fs, columns, self.schema)
 
     def __call__(self, stripe_info):
         path, stripe = stripe_info
diff --git a/dask/dataframe/io/parquet/core.py b/dask/dataframe/io/parquet/core.py
index a1d98914b..e11de5e32 100644
--- a/dask/dataframe/io/parquet/core.py
+++ b/dask/dataframe/io/parquet/core.py
@@ -1,4 +1,3 @@
-import copy
 import math
 import warnings
 from distutils.version import LooseVersion
@@ -50,7 +49,7 @@ class ParquetFunctionWrapper:
         kwargs,
         common_kwargs,
     ):
-        self.read_partition = engine.read_partition
+        self.engine = engine
         self.fs = fs
         self.meta = meta
         self.columns = columns
@@ -70,9 +69,15 @@ class ParquetFunctionWrapper:
         """
         if columns == self.columns:
             return self
-        func = copy.deepcopy(self)
-        func.columns = columns
-        return func
+        return ParquetFunctionWrapper(
+            self.engine,
+            self.fs,
+            self.meta,
+            columns,
+            self.index,
+            None,  # Already merged into common_kwargs
+            self.common_kwargs,
+        )
 
     def __call__(self, part):
 
@@ -81,7 +86,7 @@ class ParquetFunctionWrapper:
 
         return read_parquet_part(
             self.fs,
-            self.read_partition,
+            self.engine.read_partition,
             self.meta,
             [(p["piece"], p.get("kwargs", {})) for p in part],
             self.columns,


Test patch:
diff --git a/dask/dataframe/io/tests/test_hdf.py b/dask/dataframe/io/tests/test_hdf.py
index a5695de37..7122e85c1 100644
--- a/dask/dataframe/io/tests/test_hdf.py
+++ b/dask/dataframe/io/tests/test_hdf.py
@@ -9,7 +9,9 @@ import pytest
 import dask
 import dask.dataframe as dd
 from dask.dataframe._compat import tm
+from dask.dataframe.optimize import optimize_dataframe_getitem
 from dask.dataframe.utils import assert_eq
+from dask.layers import DataFrameIOLayer
 from dask.utils import dependency_depth, tmpdir, tmpfile
 
 
@@ -123,6 +125,16 @@ def test_to_hdf_multiple_nodes():
         out = dd.read_hdf(fn, "/data*")
         assert_eq(df16, out)
 
+    # Test getitem optimization
+    with tmpfile("h5") as fn:
+        a.to_hdf(fn, "/data*")
+        out = dd.read_hdf(fn, "/data*")[["x"]]
+        dsk = optimize_dataframe_getitem(out.dask, keys=out.__dask_keys__())
+        read = [key for key in dsk.layers if key.startswith("read-hdf")][0]
+        subgraph = dsk.layers[read]
+        assert isinstance(subgraph, DataFrameIOLayer)
+        assert subgraph.columns == ["x"]
+
 
 def test_to_hdf_multiple_files():
     pytest.importorskip("tables")
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
index 9fabfbfc8..434e74815 100644
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -2862,6 +2862,20 @@ def test_applymap():
     assert_eq(ddf.applymap(lambda x: (x, x)), df.applymap(lambda x: (x, x)))
 
 
+def test_add_prefix():
+    df = pd.DataFrame({"x": [1, 2, 3, 4, 5], "y": [4, 5, 6, 7, 8]})
+    ddf = dd.from_pandas(df, npartitions=2)
+    assert_eq(ddf.add_prefix("abc"), df.add_prefix("abc"))
+    assert_eq(ddf.x.add_prefix("abc"), df.x.add_prefix("abc"))
+
+
+def test_add_suffix():
+    df = pd.DataFrame({"x": [1, 2, 3, 4, 5], "y": [4, 5, 6, 7, 8]})
+    ddf = dd.from_pandas(df, npartitions=2)
+    assert_eq(ddf.add_suffix("abc"), df.add_suffix("abc"))
+    assert_eq(ddf.x.add_suffix("abc"), df.x.add_suffix("abc"))
+
+
 def test_abs():
     df = pd.DataFrame(
         {
diff --git a/dask/tests/test_distributed.py b/dask/tests/test_distributed.py
index 9925b2054..f91046676 100644
--- a/dask/tests/test_distributed.py
+++ b/dask/tests/test_distributed.py
@@ -316,11 +316,7 @@ def test_blockwise_array_creation(c, io, fuse):
 
 @pytest.mark.parametrize(
     "io",
-    [
-        "parquet-pyarrow",
-        "parquet-fastparquet",
-        "csv",
-    ],
+    ["parquet-pyarrow", "parquet-fastparquet", "csv", "hdf"],
 )
 @pytest.mark.parametrize("fuse", [True, False, None])
 def test_blockwise_dataframe_io(c, tmpdir, io, fuse):
@@ -342,6 +338,11 @@ def test_blockwise_dataframe_io(c, tmpdir, io, fuse):
     elif io == "csv":
         ddf0.to_csv(str(tmpdir), index=False)
         ddf = dd.read_csv(os.path.join(str(tmpdir), "*"))
+    elif io == "hdf":
+        pytest.importorskip("tables")
+        fn = str(tmpdir.join("h5"))
+        ddf0.to_hdf(fn, "/data*")
+        ddf = dd.read_hdf(fn, "/data*")
 
     df = df[["x"]] + 10
     ddf = ddf[["x"]] + 10