Repository: NVIDIA-Merlin/NVTabular (language: python)

Issue:
[FEA] Partition output parquet files by a column
**Is your feature request related to a problem? Please describe.**
For a realistic temporal evaluation protocol, it is common to split users interactions/sessions by time windows (e.g.  hours, days, weeks,  months), so that models can be trained over a time window and evaluated for a future time window.

So, NVTabular should be able to export the parquet files partitioned by a column (which can be a feature extracted from a timestamp column by a `LambdaOp`, like hour, day, week or month)

**Describe the solution you'd like**
When the workflow exports the parquet files, they should be organized in folders named after the partition column values (e.g. "interaction_date=2021-03-05").
That is, each folder should contain parquet files whose rows have the same value for the partition column as in their folder name
```python
workflow.transform(dataset).to_parquet('base/path', partition_by='interaction_date')
```

**Additional context**
With PySpark, this can be accomplished as follows

```python
interactions_df.write.partitionBy('session_start_date') \
                     .parquet('base/path')
```
P.s. This issue was extracted from #355  , which is broader in scope, so that it is implemented independently.
[FEA] Sequential / Session-based recommendation and time series support - Group by sorting values by timestamp
**Is your feature request related to a problem? Please describe.**
In order to support **session-based / sequence recommendation** and also **time series use cases**, we need NVTabular to be able to group data by some columns (e.g. user id, session id), sorted by another column (e.g. usually timestamp), and aggregate other columns with these different aggregation functions that takes order into account: '**list**', '**first**', '**last**'

**Describe the solution you'd like**
```python
groupby_features = ['user_id', 'session_id', 'timestamp', 'prod_id', 'prod_categ'] >> ops.Groupby(
    groupby_cols=['user_id', 'session_id'], sort_column=['timestamp'], 
    aggregated_cols={'prod_id': 'list',
                   'prod_categ': 'list',
                   'timestamp': ['first', 'last']}
)
processor = nvtabular.Workflow(groupby_features)
```

**Describe alternatives you've considered**
It might be ok is some NVT ops do not support the resulting list columns, like `Categorify` and `Standardize`, because those ops can be done before the grouping.
But the `LambdaOp` should support list columns to allow, for example, to extract the length of the lists or truncate the lists.

**Additional context**
This feature can be accomplished by different data frame frameworks:

***Pandas***
In this example, the preliminary sorting of the rows is respected by the group by, so that the aggregated columns ('item_id') will be sorted by timestamp.

```python
interactions_df.sort_values('timestamp', inplace=True)
sessions_df = interactions_df.groupby('session_id').agg('item_id': 'list')
```
P.s. `cudf / dask_cudf` also supports the 'list' aggregation function. But it does not guarantee that the data frame ordering will be respected, as pandas does.

***PySpark***
```python
import pyspark.sql.functions as F

# This window ensures the sorting of aggregated columns by another column: timestamp
session_window = Window.partitionBy('user_session').orderBy('timestamp')

df.select('user_session',                        
                        F.collect_list('product_idx').over(session_window).alias('product_seq'), 
                        F.first('event_time_ts').over(session_window).alias('session_start_ts'),
                        F.last('event_time_ts').over(session_window).alias('session_end_ts'),
                        ) \
                .groupBy('user_session').agg(
                    # The F.max() is necessary for the list column because F.collect_list() will create an incremental
                    # list of values for each value, sorted by time, and F.max() returns the longest list (containing all items of the sequence).
                    F.max('product_seq').alias('product_seq'),
                    F.min('session_start_ts').alias('session_start_ts'),
                    F.max('session_end_ts').alias('session_end_ts'),
                 )
```
This issue was extracted from #355 , which is broader in scope, so that it is implemented independently.
Other related issues: #92 #325 
[FEA] Sequential / Session-based recommendation and time series support - Group by sorting values by timestamp
**Is your feature request related to a problem? Please describe.**
In order to support **session-based / sequence recommendation** and also **time series use cases**, we need NVTabular to be able to group data by some columns (e.g. user id, session id), sorted by another column (e.g. usually timestamp), and aggregate other columns with these different aggregation functions that takes order into account: '**list**', '**first**', '**last**'

**Describe the solution you'd like**
```python
groupby_features = ['user_id', 'session_id', 'timestamp', 'prod_id', 'prod_categ'] >> ops.Groupby(
    groupby_cols=['user_id', 'session_id'], sort_column=['timestamp'], 
    aggregated_cols={'prod_id': 'list',
                   'prod_categ': 'list',
                   'timestamp': ['first', 'last']}
)
processor = nvtabular.Workflow(groupby_features)
```

**Describe alternatives you've considered**
It might be ok is some NVT ops do not support the resulting list columns, like `Categorify` and `Standardize`, because those ops can be done before the grouping.
But the `LambdaOp` should support list columns to allow, for example, to extract the length of the lists or truncate the lists.

**Additional context**
This feature can be accomplished by different data frame frameworks:

***Pandas***
In this example, the preliminary sorting of the rows is respected by the group by, so that the aggregated columns ('item_id') will be sorted by timestamp.

```python
interactions_df.sort_values('timestamp', inplace=True)
sessions_df = interactions_df.groupby('session_id').agg('item_id': 'list')
```
P.s. `cudf / dask_cudf` also supports the 'list' aggregation function. But it does not guarantee that the data frame ordering will be respected, as pandas does.

***PySpark***
```python
import pyspark.sql.functions as F

# This window ensures the sorting of aggregated columns by another column: timestamp
session_window = Window.partitionBy('user_session').orderBy('timestamp')

df.select('user_session',                        
                        F.collect_list('product_idx').over(session_window).alias('product_seq'), 
                        F.first('event_time_ts').over(session_window).alias('session_start_ts'),
                        F.last('event_time_ts').over(session_window).alias('session_end_ts'),
                        ) \
                .groupBy('user_session').agg(
                    # The F.max() is necessary for the list column because F.collect_list() will create an incremental
                    # list of values for each value, sorted by time, and F.max() returns the longest list (containing all items of the sequence).
                    F.max('product_seq').alias('product_seq'),
                    F.min('session_start_ts').alias('session_start_ts'),
                    F.max('session_end_ts').alias('session_end_ts'),
                 )
```
This issue was extracted from #355 , which is broader in scope, so that it is implemented independently.
Other related issues: #92 #325 


Discussion / hints:

This feature is essential in our time series workflows as well. In our core workflow, the core pseudocode looks like this:

```python
df = load_some_df()
df['primary_id'] = flatten_ids(['ids','to','combine'])
df.sort_values('time_id')
df.groupby('primary_id')
```
The changes proposed here would cover this case.
We also allow for per-group scaling down to the standard distribution. This would likely also be relevant to session based recommendation (scaling values by user or other key)



On another note, on the dataloading side, it will likely also be pertinent to allow for sliding-window access-time tensor building on the collected lists. As an example (ignoring NVT conventions to illustrate the point):

```python
list = [10,20,30]
session_loader = nvt.session_loader(list,window=2)
session_loader[0] #Result is [10,20]
session_loader[1] #Result is [20,30]

```
Thanks for writing up this issue @gabrielspmoreira - As you know, the **transform** `ops` available in NVTabular are all "partition-wise" operations (they do not require any shuffling of data between distinct dask.DataFrame partitions).  There are stat operators that **do** move data around, but the statistics/groupby results are always expected to be relatvely small, and the results are always persisted to disk at the end of the `fit` stage.  The `GroupbyJoin` operator, for example, is still a partition-wise transform, because the (relatively-small) groupby statistics are joined to the original data during the transform stage.

With the above in mind, I want to clarify: Is this issue requesting a **transform** operator that does not follow the "partition-wise" convention?  That is, is the goal to introduce a new type of transform that cannot be performed with a linear pass over the underlying ddf partitions or would this be a new "statistics" operation (only to be performed at "fit" time)?  In order to avoid a departure from the current/simple transform convention, it may actually make sense to separate this feature from the usual `Workflow`/`fit`/`tansform` pipeline altogether, and just introduce a `DataFrame.groupby` method.

cc @benfred (since I'm certainly interested in your thoughts here as well)
Thanks for your analysis @rjzamora. I understand that a GroupBy op requires full shuffle, because other rows from the same group (e.g. user id or session id) might not be in the same partition, right?

In this sense, I understand that it would be easier to implement `DataFrame.groupby` method outside the workflow. Therefore, it should be possible to do some custom sequence-wise transforms on the list columns after grouping. like: truncating the lists to a maximum length, counting the length of the sequence, getting the first value of the sequence. Such ops could be performed using a LambdaOp within our workflow, or our user could possibly implement then with dask_cudf `apply` outside the NVTabular workflow.

Another option would be having the GroupBy op as the first step in our pipeline, and then performing the subsequent ops (e.g. Categorify, LopOp, ...) on the list columns. 

During inference, the task of sequential recommendation and session-based recommendation is to provide a ranked list of items for a user or session, represented by a sequence of his past interactions. So, Triton should receive a batch with the last user interactions and return the recommendation list. The GroupBy op (ordered by the timestamp column), should be performed on inference the same way it was during preprocessing. So I understand that the GroupBy op should be within our workflow.
>Thanks for your analysis @rjzamora. I understand that a GroupBy op requires full shuffle, because other rows from the same group (e.g. user id or session id) might not be in the same partition, right?

Exactly right - I was using the term "partition-wise" to describe a transform that does not require a row/record to move between partitions, but I'm not sure of the best language here.

>Another option would be having the GroupBy op as the first step in our pipeline, and then performing the subsequent ops (e.g. Categorify, LopOp, ...) on the list columns.

The only reason I was suggesting a `Dataset.groupby` method (in lieue of a `Groupby` operator) is to make it impossible for users to try to use something like `ops.Groupby` in the "wrong" way.  If the grouping logic belongs to the Dataset object, and is designed to produce a **new** Dataset object, we are effectively requiring this operation to be the first step in a typical NVT pipeline.  In practice, users will get **much** better down-stream transform/dataloading performance if the result of the initial `Dataset.groupby` operation is persisted to disk (because it would relieve down-stream memory pressure), but an intermediate `to_parquet` call would not be a requirement.

  Note that my current thinking is a bit different (see below), but I am still thinking that `ops.Groupby` cannot move data between partitions.

>During inference, the task of sequential recommendation and session-based recommendation is to provide a ranked list of items for a user or session, represented by a sequence of his past interactions. So, Triton should receive a batch with the last user interactions and return the recommendation list. The GroupBy op (ordered by the timestamp column), should be performed on inference the same way it was during preprocessing. So I understand that the GroupBy op should be within our workflow.

Ah - This is a great point.  Then I agree that we ultimately **do need** something like `ops.Groupby`.  However, this operation will probably need to be a partition-wise transformation, and we will still need a new `Dataset.<something>` method to handle our data-shuffling needs for the training data.  Perhaps it would be sufficient to enable the user to shuffle the initial Dataset object by the groupby keys?  For example, if training/inference needs to group a batch of data by `['user_id', 'session_id’]`, this can be a partition-wise operation if the input data is guaranteed to have unique 'user_id’-‘session_id’ pairs within the same partition.  So, perhaps we add **both** a `Dataset.group_shuffle` method and `ops.Groupby`?  The only problem with this approach is that  the partitioning may change if the user persists the `Dataset.group_shuffle` result to disk, and then they try to read it back for processing.  We would need a new explicit option to partition by file in both the read and write (to_parquet).
This feature is essential in our time series workflows as well. In our core workflow, the core pseudocode looks like this:

```python
df = load_some_df()
df['primary_id'] = flatten_ids(['ids','to','combine'])
df.sort_values('time_id')
df.groupby('primary_id')
```
The changes proposed here would cover this case.
We also allow for per-group scaling down to the standard distribution. This would likely also be relevant to session based recommendation (scaling values by user or other key)



On another note, on the dataloading side, it will likely also be pertinent to allow for sliding-window access-time tensor building on the collected lists. As an example (ignoring NVT conventions to illustrate the point):

```python
list = [10,20,30]
session_loader = nvt.session_loader(list,window=2)
session_loader[0] #Result is [10,20]
session_loader[1] #Result is [20,30]

```
Thanks for writing up this issue @gabrielspmoreira - As you know, the **transform** `ops` available in NVTabular are all "partition-wise" operations (they do not require any shuffling of data between distinct dask.DataFrame partitions).  There are stat operators that **do** move data around, but the statistics/groupby results are always expected to be relatvely small, and the results are always persisted to disk at the end of the `fit` stage.  The `GroupbyJoin` operator, for example, is still a partition-wise transform, because the (relatively-small) groupby statistics are joined to the original data during the transform stage.

With the above in mind, I want to clarify: Is this issue requesting a **transform** operator that does not follow the "partition-wise" convention?  That is, is the goal to introduce a new type of transform that cannot be performed with a linear pass over the underlying ddf partitions or would this be a new "statistics" operation (only to be performed at "fit" time)?  In order to avoid a departure from the current/simple transform convention, it may actually make sense to separate this feature from the usual `Workflow`/`fit`/`tansform` pipeline altogether, and just introduce a `DataFrame.groupby` method.

cc @benfred (since I'm certainly interested in your thoughts here as well)
Thanks for your analysis @rjzamora. I understand that a GroupBy op requires full shuffle, because other rows from the same group (e.g. user id or session id) might not be in the same partition, right?

In this sense, I understand that it would be easier to implement `DataFrame.groupby` method outside the workflow. Therefore, it should be possible to do some custom sequence-wise transforms on the list columns after grouping. like: truncating the lists to a maximum length, counting the length of the sequence, getting the first value of the sequence. Such ops could be performed using a LambdaOp within our workflow, or our user could possibly implement then with dask_cudf `apply` outside the NVTabular workflow.

Another option would be having the GroupBy op as the first step in our pipeline, and then performing the subsequent ops (e.g. Categorify, LopOp, ...) on the list columns. 

During inference, the task of sequential recommendation and session-based recommendation is to provide a ranked list of items for a user or session, represented by a sequence of his past interactions. So, Triton should receive a batch with the last user interactions and return the recommendation list. The GroupBy op (ordered by the timestamp column), should be performed on inference the same way it was during preprocessing. So I understand that the GroupBy op should be within our workflow.
>Thanks for your analysis @rjzamora. I understand that a GroupBy op requires full shuffle, because other rows from the same group (e.g. user id or session id) might not be in the same partition, right?

Exactly right - I was using the term "partition-wise" to describe a transform that does not require a row/record to move between partitions, but I'm not sure of the best language here.

>Another option would be having the GroupBy op as the first step in our pipeline, and then performing the subsequent ops (e.g. Categorify, LopOp, ...) on the list columns.

The only reason I was suggesting a `Dataset.groupby` method (in lieue of a `Groupby` operator) is to make it impossible for users to try to use something like `ops.Groupby` in the "wrong" way.  If the grouping logic belongs to the Dataset object, and is designed to produce a **new** Dataset object, we are effectively requiring this operation to be the first step in a typical NVT pipeline.  In practice, users will get **much** better down-stream transform/dataloading performance if the result of the initial `Dataset.groupby` operation is persisted to disk (because it would relieve down-stream memory pressure), but an intermediate `to_parquet` call would not be a requirement.

  Note that my current thinking is a bit different (see below), but I am still thinking that `ops.Groupby` cannot move data between partitions.

>During inference, the task of sequential recommendation and session-based recommendation is to provide a ranked list of items for a user or session, represented by a sequence of his past interactions. So, Triton should receive a batch with the last user interactions and return the recommendation list. The GroupBy op (ordered by the timestamp column), should be performed on inference the same way it was during preprocessing. So I understand that the GroupBy op should be within our workflow.

Ah - This is a great point.  Then I agree that we ultimately **do need** something like `ops.Groupby`.  However, this operation will probably need to be a partition-wise transformation, and we will still need a new `Dataset.<something>` method to handle our data-shuffling needs for the training data.  Perhaps it would be sufficient to enable the user to shuffle the initial Dataset object by the groupby keys?  For example, if training/inference needs to group a batch of data by `['user_id', 'session_id’]`, this can be a partition-wise operation if the input data is guaranteed to have unique 'user_id’-‘session_id’ pairs within the same partition.  So, perhaps we add **both** a `Dataset.group_shuffle` method and `ops.Groupby`?  The only problem with this approach is that  the partitioning may change if the user persists the `Dataset.group_shuffle` result to disk, and then they try to read it back for processing.  We would need a new explicit option to partition by file in both the read and write (to_parquet).

Fix patch:
diff --git a/nvtabular/io/dask.py b/nvtabular/io/dask.py
index 4ae9b12c7aa..00a7d13800f 100644
--- a/nvtabular/io/dask.py
+++ b/nvtabular/io/dask.py
@@ -18,6 +18,7 @@
 import dask
 import pandas as pd
 from dask.base import tokenize
+from dask.dataframe.core import _concat, new_dd_object
 from dask.delayed import Delayed
 from dask.highlevelgraph import HighLevelGraph
 from nvtx import annotate
@@ -94,6 +95,78 @@ def _write_output_partition(
     return df_size
 
 
+def _get_partition_groups(df, partition_cols, fs, output_path, filename):
+    # Logic copied from cudf/cudf/io/parquet.py
+    df = df.sort_values(partition_cols)
+    divisions = df[partition_cols].drop_duplicates(ignore_index=True)
+    splits = df[partition_cols].searchsorted(divisions, side="left")
+    splits = splits.tolist() + [len(df[partition_cols])]
+
+    subgroups, fns = [], []
+    for i in range(0, len(splits) - 1):
+        sub_df = df.iloc[splits[i] : splits[i + 1]].copy(deep=False)
+        if sub_df is None or len(sub_df) == 0:
+            continue
+        keys = tuple([sub_df[col].iloc[0] for col in partition_cols])
+        if not isinstance(keys, tuple):
+            keys = (keys,)
+        subdir = fs.sep.join(
+            [
+                "{colname}={value}".format(colname=name, value=val)
+                for name, val in zip(partition_cols, keys)
+            ]
+        )
+        prefix = fs.sep.join([output_path, subdir])
+        fs.mkdirs(prefix, exist_ok=True)
+        fns.append(fs.sep.join([subdir, filename]))
+
+        write_df = sub_df.copy(deep=False)
+        write_df.drop(columns=partition_cols, inplace=True)
+        subgroups.append(write_df)
+
+    return fns, subgroups
+
+
+@annotate("write_partitioned", color="green", domain="nvt_python")
+def _write_partitioned(
+    df,
+    filename,
+    output_path,
+    partition_cols,
+    shuffle,
+    fs,
+    cat_names,
+    cont_names,
+    label_names,
+    output_format,
+    num_threads,
+    cpu,
+):
+    # Logic copied from cudf/cudf/io/parquet.py
+    data_cols = df.columns.drop(partition_cols)
+    if len(data_cols) == 0:
+        raise ValueError("No data left to save outside partition columns")
+
+    # Partition the input data
+    fns, dfs = _get_partition_groups(df, partition_cols, fs, output_path, filename)
+    writer = writer_factory(
+        output_format,
+        output_path,
+        1,
+        shuffle,
+        num_threads=num_threads,
+        cpu=cpu,
+        fns=fns,
+    )
+    writer.set_col_names(labels=label_names, cats=cat_names, conts=cont_names)
+
+    # Add the partitioned data
+    writer.add_data(dfs)
+
+    # Return metadata and row-count in dict
+    return writer.close()
+
+
 @annotate("write_subgraph", color="green", domain="nvt_python")
 def _write_subgraph(
     subgraph,
@@ -153,6 +226,23 @@ def _write_metadata_files(md_list, output_path, output_format, cpu):
     wc.write_special_metadata(special_md, fs, output_path)
 
 
+def _simple_shuffle(ddf, plan):
+
+    # Construct graph for a simple shuffle
+    token = tokenize(ddf, plan)
+    name = "shuffled-" + token
+    final_tasks = collections.defaultdict(list)
+    ignore_index = True
+    for i, p in enumerate(plan):
+        final_tasks[(name, p)].append((ddf._name, i))
+    dsk = {k: (_concat, v, ignore_index) for k, v in final_tasks.items()}
+
+    # Conver to a DataFrame collection
+    graph = HighLevelGraph.from_collections(name, dsk, dependencies=[ddf])
+    divisions = [None] * (len(dsk) + 1)
+    return new_dd_object(graph, name, ddf._meta, divisions)
+
+
 def _ddf_to_dataset(
     ddf,
     fs,
@@ -168,10 +258,13 @@ def _ddf_to_dataset(
     num_threads,
     cpu,
     suffix="",
+    partition_on=None,
 ):
 
     # Construct graph for Dask-based dataset write
-    token = tokenize(ddf, shuffle, out_files_per_proc, cat_names, cont_names, label_names)
+    token = tokenize(
+        ddf, shuffle, out_files_per_proc, cat_names, cont_names, label_names, suffix, partition_on
+    )
     name = "write-processed-" + token
     write_name = name + "-partition" + token
 
@@ -180,7 +273,34 @@ def _ddf_to_dataset(
 
     dsk = {}
     task_list = []
-    if file_partition_map is not None:
+    if partition_on:
+        # Use hive partitioning to write the data
+        cached_writers = False
+        for idx in range(ddf.npartitions):
+            task_list.append((write_name, idx))
+            dsk[task_list[-1]] = (
+                _write_partitioned,
+                (ddf._name, idx),
+                f"part.{idx}{suffix}",
+                output_path,
+                partition_on,
+                shuffle,
+                fs,
+                cat_names,
+                cont_names,
+                label_names,
+                output_format,
+                num_threads,
+                cpu,
+            )
+        dsk[name] = (
+            _write_metadata_files,
+            task_list,
+            output_path,
+            output_format,
+            cpu,
+        )
+    elif file_partition_map is not None:
         # Use specified mapping of data to output files
         cached_writers = False
         full_graph = ddf.dask
diff --git a/nvtabular/io/dataset.py b/nvtabular/io/dataset.py
index 05387012e13..0f3d7b1cd9a 100644
--- a/nvtabular/io/dataset.py
+++ b/nvtabular/io/dataset.py
@@ -14,6 +14,7 @@
 # limitations under the License.
 #
 
+import collections
 import logging
 import math
 import random
@@ -36,7 +37,7 @@
 
 from ..utils import device_mem_size
 from .csv import CSVDatasetEngine
-from .dask import _ddf_to_dataset
+from .dask import _ddf_to_dataset, _simple_shuffle
 from .dataframe_engine import DataFrameDatasetEngine
 from .parquet import ParquetDatasetEngine
 
@@ -373,6 +374,109 @@ def to_gpu(self):
         self.cpu = False
         self.engine.to_gpu()
 
+    def shuffle_by_keys(self, keys, hive_data=None, npartitions=None):
+        """Shuffle the in-memory Dataset so that all unique-key
+        combinations are moved to the same partition.
+
+        Parameters
+        ----------
+        keys : list(str)
+            Column names to shuffle by.
+        hive_data : bool; default None
+            Whether the dataset is backed by a hive-partitioned
+            dataset (with the keys encoded in the directory structure).
+            By default, the Dataset's `file_partition_map` property will
+            be inspected to infer this setting. When `hive_data` is True,
+            the number of output partitions will correspond to the number
+            of unique key combinations in the dataset.
+        npartitions : int; default None
+            Number of partitions in the output Dataset. For hive-partitioned
+            data, this value should be <= the number of unique key
+            combinations (the default), otherwise it will be ignored. For
+            data that is not hive-partitioned, the ``npartitions`` input
+            should be <= the orginal partition count, otherwise it will be
+            ignored.
+        """
+
+        # Make sure we are dealing with a list
+        keys = [keys] if not isinstance(keys, (list, tuple)) else keys
+
+        # Start with default ddf
+        ddf = self.to_ddf()
+        if npartitions:
+            npartitions = min(ddf.npartitions, npartitions)
+
+        if hive_data is not False:
+            # The keys may be encoded in the directory names.
+            # Let's use the file_partition_map to extract this info.
+            try:
+                _mapping = self.file_partition_map
+            except (AttributeError):
+                _mapping = None
+                if hive_data:
+                    raise RuntimeError("Failed to extract hive-partition mapping!")
+
+            # If we have a `_mapping` available, check if the
+            # file names include information about all our keys
+            hive_mapping = collections.defaultdict(list)
+            if _mapping:
+                for k, v in _mapping.items():
+                    for part in k.split(self.engine.fs.sep)[:-1]:
+                        try:
+                            _key, _val = part.split("=")
+                        except ValueError:
+                            continue
+                        if _key in keys:
+                            hive_mapping[_key].append(_val)
+
+            if set(hive_mapping.keys()) == set(keys):
+
+                # Generate hive-mapping DataFrame summary
+                hive_mapping = type(ddf._meta)(hive_mapping)
+                cols = list(hive_mapping.columns)
+                for c in keys:
+                    typ = ddf._meta[c].dtype
+                    if c in cols:
+                        hive_mapping[c] = hive_mapping[c].astype(typ)
+
+                # Generate simple-shuffle plan
+                target_mapping = hive_mapping.drop_duplicates().reset_index(drop=True)
+                target_mapping.index.name = "_partition"
+                hive_mapping.index.name = "_sort"
+                target_mapping.reset_index(drop=False, inplace=True)
+                plan = (
+                    hive_mapping.reset_index()
+                    .merge(target_mapping, on=cols, how="left")
+                    .sort_values("_sort")["_partition"]
+                )
+
+                if hasattr(plan, "to_pandas"):
+                    plan = plan.to_pandas()
+
+                # Deal with repartitioning
+                if npartitions and npartitions < len(target_mapping):
+                    q = np.linspace(0.0, 1.0, num=npartitions + 1)
+                    divs = plan.quantile(q)
+                    partitions = divs.searchsorted(plan, side="right") - 1
+                    partitions[(plan >= divs.iloc[-1]).values] = len(divs) - 2
+                    plan = partitions.tolist()
+                elif len(plan) != len(plan.unique()):
+                    plan = plan.to_list()
+                else:
+                    # Plan is a unique 1:1 ddf partition mapping.
+                    # We already have shuffled data.
+                    return self
+
+                # TODO: We should avoid shuffling the original ddf and
+                # instead construct a new (more-efficent) graph to read
+                # multiple files from each partition directory at once.
+                # Generally speaking, we can optimize this code path
+                # much further.
+                return Dataset(_simple_shuffle(ddf, plan))
+
+        # Fall back to dask.dataframe algorithm
+        return Dataset(ddf.shuffle(keys, npartitions=npartitions))
+
     def to_iter(self, columns=None, indices=None, shuffle=False, seed=None):
         """Convert `Dataset` object to a `cudf.DataFrame` iterator.
 
@@ -419,6 +523,7 @@ def to_parquet(
         conts=None,
         labels=None,
         suffix=".parquet",
+        partition_on=None,
     ):
         """Writes out to a parquet dataset
 
@@ -427,8 +532,7 @@ def to_parquet(
         output_path : string
             Path to write processed/shuffled output data
         shuffle : nvt.io.Shuffle enum
-            How to shuffle the output dataset. Shuffling is only
-            performed if the data is written to disk. For all options,
+            How to shuffle the output dataset. For all options,
             other than `None` (which means no shuffling), the partitions
             of the underlying dataset/ddf will be randomly ordered. If
             `PER_PARTITION` is specified, each worker/process will also
@@ -441,6 +545,11 @@ def to_parquet(
             data processed by each worker.  To improve performace, this option
             currently uses host-memory `BytesIO` objects for the intermediate
             persist stage. The `FULL` option is not yet implemented.
+        partition_on : str or list(str)
+            Columns to use for hive-partitioning.  If this option is used,
+            `preserve_files`, `output_files`, and `out_files_per_proc` will
+            all be ignored.  Also, the `PER_WORKER` shuffle will not be
+            supported.
         preserve_files : bool
             Whether to preserve the original file-to-partition mapping of
             the base dataset. This option is only available if the base
@@ -556,6 +665,7 @@ def to_parquet(
             num_threads,
             self.cpu,
             suffix=suffix,
+            partition_on=partition_on,
         )
 
     def to_hugectr(
diff --git a/nvtabular/io/parquet.py b/nvtabular/io/parquet.py
index ac084732554..e7885f17bfa 100644
--- a/nvtabular/io/parquet.py
+++ b/nvtabular/io/parquet.py
@@ -782,8 +782,8 @@ def write_special_metadata(cls, md, fs, out_dir):
 
     def _close_writers(self):
         md_dict = {}
-        for writer, path in zip(self.data_writers, self.data_paths):
-            fn = path.split(self.fs.sep)[-1]
+        _fns = self.fns or [path.split(self.fs.sep)[-1] for path in self.data_paths]
+        for writer, fn in zip(self.data_writers, _fns):
             md_dict[fn] = writer.close(metadata_file_path=fn)
         for f in self.data_files:
             f.close()
@@ -853,10 +853,11 @@ def write_special_metadata(cls, md, fs, out_dir):
         _write_pq_metadata_file_pyarrow(md_list, fs, out_dir)
 
     def _close_writers(self):
-        for writer, path in zip(self.data_writers, self.data_paths):
-            fn = path.split(self.fs.sep)[-1]
+        _fns = self.fns or [path.split(self.fs.sep)[-1] for path in self.data_paths]
+        for writer, fn in zip(self.data_writers, _fns):
             writer.close()
-            self.md_collectors[path][0].set_file_path(fn)
+            _path = self.fs.sep.join([self.out_dir, fn])
+            self.md_collectors[_path][0].set_file_path(fn)
         return self.md_collectors
 
 
diff --git a/nvtabular/io/writer.py b/nvtabular/io/writer.py
index ab4d9e01c42..c9889b86996 100644
--- a/nvtabular/io/writer.py
+++ b/nvtabular/io/writer.py
@@ -85,7 +85,7 @@ def __init__(
             self.num_out_files = len(self.fns)
         else:
             self.num_out_files = num_out_files
-        self.num_samples = [0] * num_out_files
+        self.num_samples = [0] * self.num_out_files
 
         self.data_paths = None
         self.need_cal_col_names = True
@@ -128,10 +128,14 @@ def _write_thread(self):
     def add_data(self, df):
         # Populate columns idxs
         if not self.col_idx:
-            for i, x in enumerate(df.columns.values):
+            _df = df[0] if isinstance(df, list) else df
+            for i, x in enumerate(_df.columns.values):
                 self.col_idx[str(x)] = i
 
-        if self.num_out_files == 1:
+        if isinstance(df, list):
+            # Data is already partitioned into a list
+            self._add_data_list(df)
+        elif self.num_out_files == 1:
             # Only writing to a single file. No need to
             # scatter or slice the data before writing
             self._add_single_file(df)
@@ -199,6 +203,18 @@ def _add_data_slice(self, df):
             else:
                 self._write_table(x, to_write)
 
+    def _add_data_list(self, dfs):
+        """Write a list of DataFrames"""
+        assert len(dfs) == self.num_out_files
+        for x, df in enumerate(dfs):
+            if self.shuffle:
+                df = _shuffle_df(df)
+            self.num_samples[x] = self.num_samples[x] + df.shape[0]
+            if self.num_threads > 1:
+                self.queue.put((x, df))
+            else:
+                self._write_table(x, df)
+
     def _add_single_file(self, df):
         """Write to single file.
 


Test patch:
diff --git a/tests/unit/test_io.py b/tests/unit/test_io.py
index 4a1c9294240..40d0226dbcb 100644
--- a/tests/unit/test_io.py
+++ b/tests/unit/test_io.py
@@ -601,3 +601,105 @@ def test_dataset_conversion(tmpdir, cpu, preserve_files):
     # Check that the `suffix=".pq"` argument was successful
     assert glob.glob(os.path.join(pq_path, "*.pq"))
     assert not glob.glob(os.path.join(pq_path, "*.parquet"))
+
+
+@pytest.mark.parametrize("cpu", [True, False])
+def test_hive_partitioned_data(tmpdir, cpu):
+
+    # Initial timeseries dataset (in cpu memory).
+    # Round the full "timestamp" to the hour for partitioning.
+    ddf = dask.datasets.timeseries(
+        start="2000-01-01",
+        end="2000-01-03",
+        freq="600s",
+        partition_freq="6h",
+        seed=42,
+    ).reset_index()
+    ddf["timestamp"] = ddf["timestamp"].dt.round("D").dt.day
+    ds = nvt.Dataset(ddf, engine="parquet")
+
+    # Write the dataset to disk
+    path = str(tmpdir)
+    partition_keys = ["timestamp", "name"]
+    ds.to_parquet(path, partition_on=partition_keys)
+
+    # Make sure the directory structure is hive-like
+    df_expect = ddf.compute()
+    df_expect = df_expect.sort_values(["id", "x", "y"]).reset_index(drop=True)
+    timestamp_check = df_expect["timestamp"].iloc[0]
+    name_check = df_expect["name"].iloc[0]
+    assert glob.glob(
+        os.path.join(
+            path,
+            f"timestamp={timestamp_check}/name={name_check}/*",
+        )
+    )
+
+    # Read back with dask.dataframe and check the data
+    df_check = dd.read_parquet(path).compute()
+    df_check["name"] = df_check["name"].astype("object")
+    df_check["timestamp"] = df_check["timestamp"].astype("int64")
+    df_check = df_check.sort_values(["id", "x", "y"]).reset_index(drop=True)
+    for col in df_expect:
+        # Order of columns can change after round-trip partitioning
+        assert_eq(df_expect[col], df_check[col], check_index=False)
+
+    # Read back with NVT and check the data
+    df_check = nvt.Dataset(path, engine="parquet").to_ddf().compute()
+    df_check["name"] = df_check["name"].astype("object")
+    df_check["timestamp"] = df_check["timestamp"].astype("int64")
+    df_check = df_check.sort_values(["id", "x", "y"]).reset_index(drop=True)
+    for col in df_expect:
+        # Order of columns can change after round-trip partitioning
+        assert_eq(df_expect[col], df_check[col], check_index=False)
+
+
+@pytest.mark.parametrize("cpu", [True, False])
+@pytest.mark.parametrize("partition_on", [None, ["name", "id"], ["name"]])
+@pytest.mark.parametrize("keys", [["name"], ["id"], ["name", "id"]])
+@pytest.mark.parametrize("npartitions", [None, 2])
+def test_dataset_shuffle_on_keys(tmpdir, cpu, partition_on, keys, npartitions):
+
+    # Initial timeseries dataset
+    size = 60
+    df1 = pd.DataFrame(
+        {
+            "name": np.random.choice(["Dave", "Zelda"], size=size),
+            "id": np.random.choice([0, 1], size=size),
+            "x": np.random.uniform(low=0.0, high=10.0, size=size),
+            "y": np.random.uniform(low=0.0, high=10.0, size=size),
+        }
+    )
+    ddf1 = dd.from_pandas(df1, npartitions=3)
+
+    # Write the dataset to disk
+    path = str(tmpdir)
+    ddf1.to_parquet(str(tmpdir), partition_on=partition_on)
+
+    # Construct NVT Dataset
+    ds = nvt.Dataset(path, engine="parquet")
+
+    # Shuffle the dataset by `keys`
+    ds2 = ds.shuffle_by_keys(keys, npartitions=npartitions)
+
+    # Inspect the result
+    ddf2 = ds2.to_ddf()
+    if npartitions:
+        assert ddf2.npartitions == npartitions
+
+    # A successful shuffle will return the same unique-value
+    # count for both the full dask algorithm and a partition-wise sum
+    n1 = sum([len(p[keys].drop_duplicates()) for p in ddf2.partitions])
+    n2 = len(ddf2[keys].drop_duplicates())
+    assert n1 == n2
+
+    # Check that none of the rows was changed
+    df1 = df1.sort_values(["id", "x", "y"]).reset_index(drop=True)
+    df2 = ddf2.compute().sort_values(["id", "x", "y"]).reset_index(drop=True)
+    if partition_on:
+        # Dask will convert partitioned columns to Categorical
+        df2["name"] = df2["name"].astype("object")
+        df2["id"] = df2["id"].astype("int64")
+    for col in df1:
+        # Order of columns can change after round-trip partitioning
+        assert_eq(df1[col], df2[col], check_index=False)