Skip to content

Commit d5d5eb6

Browse files
lmeyerovclaude
andcommitted
fix(tests): cuDF compatibility for tolist() calls in chain optimization tests
Use to_arrow().to_pylist() for cuDF with fallback to tolist() for pandas. Fixes test_same_nodes_with_and_without_where and test_same_edges_with_and_without_where. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
1 parent a730205 commit d5d5eb6

File tree

1 file changed

+14
-4
lines changed

1 file changed

+14
-4
lines changed

tests/gfql/ref/test_chain_optimizations.py

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -920,8 +920,13 @@ def test_same_nodes_with_and_without_where(self, linear_graph):
920920
chain_with_where = Chain(ops, where=where)
921921
result_with_where = linear_graph.gfql(chain_with_where)
922922

923-
nodes_no_where = set(result_no_where._nodes['id'].tolist())
924-
nodes_with_where = set(result_with_where._nodes['id'].tolist())
923+
# Use to_arrow().to_pylist() for cuDF compatibility
924+
try:
925+
nodes_no_where = set(result_no_where._nodes['id'].to_arrow().to_pylist())
926+
nodes_with_where = set(result_with_where._nodes['id'].to_arrow().to_pylist())
927+
except AttributeError:
928+
nodes_no_where = set(result_no_where._nodes['id'].tolist())
929+
nodes_with_where = set(result_with_where._nodes['id'].tolist())
925930

926931
assert nodes_no_where == nodes_with_where
927932

@@ -939,8 +944,13 @@ def test_same_edges_with_and_without_where(self, linear_graph):
939944
chain_with_where = Chain(ops, where=where)
940945
result_with_where = linear_graph.gfql(chain_with_where)
941946

942-
edges_no_where = set(result_no_where._edges['eid'].tolist())
943-
edges_with_where = set(result_with_where._edges['eid'].tolist())
947+
# Use to_arrow().to_pylist() for cuDF compatibility
948+
try:
949+
edges_no_where = set(result_no_where._edges['eid'].to_arrow().to_pylist())
950+
edges_with_where = set(result_with_where._edges['eid'].to_arrow().to_pylist())
951+
except AttributeError:
952+
edges_no_where = set(result_no_where._edges['eid'].tolist())
953+
edges_with_where = set(result_with_where._edges['eid'].tolist())
944954

945955
assert edges_no_where == edges_with_where
946956

0 commit comments

Comments
 (0)