Update test_partitioning.py

rapidsai · Sep 25, 2024 · 67f01ce · 67f01ce
1 parent eda42a6
commit 67f01ce
Showing 1 changed file with 8 additions and 7 deletions.
diff --git a/python/pylibcudf/pylibcudf/tests/test_partitioning.py b/python/pylibcudf/pylibcudf/tests/test_partitioning.py
@@ -10,13 +10,14 @@
 def partitioning_data():
     data = {"a": [1, 2, 3], "b": [1, 2, 5], "c": [1, 2, 10]}
     pa_table = pa.table(data)
-    return data, pa_table
+    plc_table = plc.interop.from_arrow(pa_table)
+    return data, plc_table, pa_table
 
 
 def test_partition(partitioning_data):
-    raw_data, pa_table = partitioning_data
+    raw_data, plc_table, pa_table = partitioning_data
     result, result_offsets = plc.partitioning.partition(
-        plc.interop.from_arrow(pa_table),
+        plc_table,
         plc.interop.from_arrow(pa.array([0, 0, 0])),
         1,
     )
@@ -29,9 +30,9 @@ def test_partition(partitioning_data):
 
 
 def test_hash_partition(partitioning_data):
-    raw_data, pa_table = partitioning_data
+    raw_data, plc_table, pa_table = partitioning_data
     result, result_offsets = plc.partitioning.hash_partition(
-        plc.interop.from_arrow(pa_table), [0, 1], 1
+        plc_table, [0, 1], 1
     )
     expected = pa.table(
         list(raw_data.values()),
@@ -42,9 +43,9 @@ def test_hash_partition(partitioning_data):
 
 
 def test_round_robin_partition(partitioning_data):
-    raw_data, pa_table = partitioning_data
+    raw_data, plc_table, pa_table = partitioning_data
     result, result_offsets = plc.partitioning.round_robin_partition(
-        plc.interop.from_arrow(pa_table), 1, 0
+        plc_table, 1, 0
     )
     expected = pa.table(
         list(raw_data.values()),