From 67f01cedfec0d8ff43af56b2bb20627022333f93 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Wed, 25 Sep 2024 15:57:56 -0700 Subject: [PATCH] Update test_partitioning.py --- .../pylibcudf/tests/test_partitioning.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/python/pylibcudf/pylibcudf/tests/test_partitioning.py b/python/pylibcudf/pylibcudf/tests/test_partitioning.py index 2e4f33a1bff..444d0089d2c 100644 --- 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()),