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

Issue:
[FEA] Truncate List columns (sparse tensors) - related to the GroupBy op
**Is your feature request related to a problem? Please describe.**
This feature is related with #641 "Sequential / Session-based recommendation and time series support - Group by sorting values by timestamp".
After grouping, some sequences (e.g. user sessions or time series) might be very long, and for some ML models sequences with maximum (fixed) length are required. So lists truncation is necessary.
As currently List columns are internally represented as sparse vectors, it is currently not possible to  use a LambdaOp to truncate the list values to a maximum length.

**Describe the solution you'd like**
I would like either an option on Groupby op to truncate all aggregated list columns to the same maximum length or an independent TruncateList op that would truncate selected list columns

**Describe alternatives you've considered**
As a workaround for this problem I am extending the NVT PyTorch dataloader, converting the internal representation of list columns to PyTorch sparse tensors (as shown in #500), converting them to dense tensors (with padding zeros in the right) and then slicing the second dimension of the tensor to the maximum length.
But storing longer sequences than needed in the parquet files is a waste of space and requires workarounds like this in the model size.




Discussion / hints:
This is similar to string slicing with string columns - we really need to have list slicing for list columns. 

We should prototype this with the offset/values and followup with the cudf team
There should be possible to truncate either the start (positive number) or the end of the sessions (negative number).
A first-order solution is probably to use the [list.take](https://github.com/rapidsai/cudf/blob/d56428abfc7345e2f7be3b679f2383e0c1eb7084/python/cudf/cudf/core/column/lists.py#L317) method for cudf list columns.  That is, you can just make a column of list indices, and call `take` to get what you want.  

It should be possible to get better performance with either a cudf primitive or with other custom solutions, but a quick/simple solution may be a good start.

Fix patch:
diff --git a/docs/source/resources/api/ops/index.rst b/docs/source/resources/api/ops/index.rst
index ca4d1107f1d..81a512348a2 100644
--- a/docs/source/resources/api/ops/index.rst
+++ b/docs/source/resources/api/ops/index.rst
@@ -19,6 +19,7 @@ Operators
    JoinExternal <joinexternal>
    JoinGroupby <joingroupby>
    LambdaOp <lambdaop>
+   ListSlice <listslice>
    LogOp <log>
    Normalize <normalize>
    NormalizeMinMax <normalizeminmax>
diff --git a/docs/source/resources/api/ops/listslice.rst b/docs/source/resources/api/ops/listslice.rst
new file mode 100644
index 00000000000..9e39a214b82
--- /dev/null
+++ b/docs/source/resources/api/ops/listslice.rst
@@ -0,0 +1,6 @@
+ListSlice
+=========
+
+.. autoclass:: nvtabular.ops.ListSlice
+   :members:
+   :show-inheritance:
diff --git a/nvtabular/ops/__init__.py b/nvtabular/ops/__init__.py
index ac58f51352a..2d947560a1a 100644
--- a/nvtabular/ops/__init__.py
+++ b/nvtabular/ops/__init__.py
@@ -30,6 +30,7 @@
 from .join_external import JoinExternal
 from .join_groupby import JoinGroupby
 from .lambdaop import LambdaOp
+from .list_slice import ListSlice
 from .logop import LogOp
 from .normalize import Normalize, NormalizeMinMax
 from .operator import Operator
diff --git a/nvtabular/ops/list_slice.py b/nvtabular/ops/list_slice.py
new file mode 100644
index 00000000000..cd6af21d4c2
--- /dev/null
+++ b/nvtabular/ops/list_slice.py
@@ -0,0 +1,138 @@
+#
+# Copyright (c) 2021, 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 cudf
+import cupy as cp
+import numba.cuda
+import numpy as np
+import pandas as pd
+from cudf.core.column import as_column, build_column
+from nvtx import annotate
+
+from nvtabular.dispatch import DataFrameType
+
+from .operator import ColumnNames, Operator
+
+
+class ListSlice(Operator):
+    """Slices a list column
+
+    This operator provides the ability to slice list column by row. For example, to truncate a
+    list column to only include the first 10 elements per row::
+
+        truncated = column_names >> ops.ListSlice(10)
+
+    Take the first 10 items, ignoring the first element::
+
+        truncated = column_names >> ops.ListSlice(1, 11)
+
+    Take the last 10 items from each row::
+
+        truncated = column_names >> ops.ListSlice(-10)
+    """
+
+    def __init__(self, start, end=None):
+        super().__init__()
+        self.start = start
+        self.end = end
+
+        if self.start > 0 and self.end is None:
+            self.end = self.start
+            self.start = 0
+
+        if self.end is None:
+            self.end = np.iinfo(np.int64).max
+
+    @annotate("ListSlice_op", color="darkgreen", domain="nvt_python")
+    def transform(self, columns: ColumnNames, df: DataFrameType) -> DataFrameType:
+        on_cpu = isinstance(df, pd.DataFrame)
+        ret = pd.DataFrame() if on_cpu else cudf.DataFrame()
+        for col in columns:
+            # handle CPU via normal python slicing (not very efficient)
+            if on_cpu:
+                ret[col] = [row[self.start : self.end] for row in df[col]]
+            else:
+                # figure out the size of each row from the list offsets
+                c = df[col]._column
+                offsets = c.offsets.values
+                elements = c.elements.values
+
+                # figure out the size of each row after slicing start/end
+                new_offsets = cp.zeros(offsets.size, dtype=offsets.dtype)
+                threads = 32
+                blocks = (offsets.size + threads - 1) // threads
+
+                # calculate new row offsets after slicing
+                _calculate_row_sizes[blocks, threads](self.start, self.end, offsets, new_offsets)
+                new_offsets = cp.cumsum(new_offsets).astype(offsets.dtype)
+
+                # create a new array for the sliced elements
+                new_elements = cp.zeros(new_offsets[-1].item(), dtype=elements.dtype)
+                if new_elements.size:
+                    _slice_rows[blocks, threads](
+                        self.start, offsets, elements, new_offsets, new_elements
+                    )
+
+                # build up a list column with the sliced values
+                ret[col] = build_column(
+                    None,
+                    dtype=cudf.core.dtypes.ListDtype(new_elements.dtype),
+                    size=new_offsets.size - 1,
+                    children=(as_column(new_offsets), as_column(new_elements)),
+                )
+
+        return ret
+
+    transform.__doc__ = Operator.transform.__doc__
+
+
+@numba.cuda.jit
+def _calculate_row_sizes(start, end, offsets, row_sizes):
+    """given a slice (start/end) and existing offsets indicating row lengths, this
+    calculates the size for each new row after slicing"""
+    rowid = numba.cuda.grid(1)
+    if rowid < offsets.size - 1:
+        original_row_size = offsets[rowid + 1] - offsets[rowid]
+
+        # handle negative slicing appropiately
+        if start < 0:
+            start = original_row_size + start
+        if end < 0:
+            end = original_row_size + end
+
+        # clamp start/end to be in (0, original_row_size)
+        start = min(max(0, start), original_row_size)
+        end = min(max(0, end), original_row_size)
+
+        row_sizes[rowid + 1] = end - start
+
+
+@numba.cuda.jit
+def _slice_rows(start, offsets, elements, new_offsets, new_elements):
+    """slices rows of a list column. requires the 'new_offsets' to
+    be previously calculated (meaning that we don't need the 'end' slice index
+    since thats baked into the new_offsets"""
+    rowid = numba.cuda.grid(1)
+    if rowid < (new_offsets.size - 1):
+        if start >= 0:
+            offset = offsets[rowid] + start
+        else:
+            offset = offsets[rowid + 1] + start
+            if offset < offsets[rowid]:
+                offset = offsets[rowid]
+
+        for new_offset in range(new_offsets[rowid], new_offsets[rowid + 1]):
+            new_elements[new_offset] = elements[offset]
+            offset += 1


Test patch:
diff --git a/tests/unit/test_ops.py b/tests/unit/test_ops.py
index 4fe159012bd..f262e5231fc 100644
--- a/tests/unit/test_ops.py
+++ b/tests/unit/test_ops.py
@@ -967,7 +967,6 @@ def test_data_stats(tmpdir, df, datasets, engine):
 @pytest.mark.parametrize("cpu", [False, True])
 @pytest.mark.parametrize("keys", [["name"], "id", ["name", "id"]])
 def test_groupby_op(keys, cpu):
-
     # Initial timeseries dataset
     size = 60
     df1 = pd.DataFrame(
@@ -1017,3 +1016,42 @@ def test_groupby_op(keys, cpu):
 
     # Check basic behavior or "y" column
     assert (new_gdf["y-first"] < new_gdf["y-last"]).all()
+
+
+@pytest.mark.parametrize("cpu", [True, False])
+def test_list_slice(cpu):
+    DataFrame = pd.DataFrame if cpu else cudf.DataFrame
+
+    df = DataFrame({"y": [[0, 1, 2, 2, 767], [1, 2, 2, 3], [1, 223, 4]]})
+
+    op = ops.ListSlice(0, 2)
+    print("df", df)
+    transformed = op.transform(["y"], df)
+    expected = DataFrame({"y": [[0, 1], [1, 2], [1, 223]]})
+    assert_eq(transformed, expected)
+
+    op = ops.ListSlice(3, 5)
+    print("df", df)
+    transformed = op.transform(["y"], df)
+    expected = DataFrame({"y": [[2, 767], [3], []]})
+    assert_eq(transformed, expected)
+
+    op = ops.ListSlice(4, 10)
+    transformed = op.transform(["y"], df)
+    expected = DataFrame({"y": [[767], [], []]})
+    assert_eq(transformed, expected)
+
+    op = ops.ListSlice(100, 20000)
+    transformed = op.transform(["y"], df)
+    expected = DataFrame({"y": [[], [], []]})
+    assert_eq(transformed, expected)
+
+    op = ops.ListSlice(-4)
+    transformed = op.transform(["y"], df)
+    expected = DataFrame({"y": [[1, 2, 2, 767], [1, 2, 2, 3], [1, 223, 4]]})
+    assert_eq(transformed, expected)
+
+    op = ops.ListSlice(-3, -1)
+    transformed = op.transform(["y"], df)
+    expected = DataFrame({"y": [[2, 2], [2, 2], [1, 223]]})
+    assert_eq(transformed, expected)