Use row lengths instead of offsets in convert_df_to_dict (#285)

NVIDIA-Merlin · Feb 2, 2023 · 3510a44 · 3510a44
1 parent 2be1fa4
commit 3510a44
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/merlin/systems/triton/__init__.py b/merlin/systems/triton/__init__.py
@@ -71,7 +71,9 @@ def _convert_df_to_dict(schema, batch, dtype="int32"):
                 df_dict[col_name + "__values"] = col.list.leaves.values_host.astype(
                     col_schema.dtype.to_numpy
                 )
-                df_dict[col_name + "__lengths"] = col._column.offsets.values_host.astype(dtype)
+                offsets = col._column.offsets.values_host.astype(dtype)
+                row_lengths = offsets[1:] - offsets[:-1]
+                df_dict[col_name + "__lengths"] = row_lengths
             else:
                 values = col.list.leaves.values_host
                 values = values.reshape(*shape).astype(col_schema.dtype.to_numpy)