Add doc string and refactor test

rapidsai · Jun 26, 2024 · 5de9f2b · 5de9f2b
1 parent 351d723
commit 5de9f2b
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 5 deletions.
diff --git a/python/cudf/cudf/_lib/pylibcudf/lists.pyx b/python/cudf/cudf/_lib/pylibcudf/lists.pyx
@@ -103,10 +103,20 @@ cpdef Column concatenate_list_elements(Column input, bool dropna):
 
 
 cpdef Column count_elements(Column input):
-    # shared_ptr required because lists_column_view has no default
-    # ctor
+    """Count the number of rows in each list element in the given lists column.
+
+    Parameters
+    ----------
+    input : Column
+        The input column
+
+    Returns
+    -------
+    Column
+        A new Column of the lengths of each list element
+    """
     cdef shared_ptr[lists_column_view] list_view = (
-        make_shared[lists_column_view](col.view())
+        make_shared[lists_column_view](input.view())
     )
     cdef unique_ptr[column] c_result
 

diff --git a/python/cudf/cudf/pylibcudf_tests/test_lists.py b/python/cudf/cudf/pylibcudf_tests/test_lists.py
@@ -50,8 +50,8 @@ def test_count_elements():
     test_data = [[1, 2, 3], [4], [5, 6]]
     arr = pa.array(test_data)
     plc_column = plc.interop.from_arrow(arr)
-    res = plc.lists.contains(plc_column)
+    res = plc.lists.count_elements(plc_column)
 
-    expect = pa.array([3, 1, 2])
+    expect = pa.array([3, 1, 2], type=pa.int32())
 
     assert_column_eq(expect, res)