Skip to content

Commit 607f42c

Browse files
lmeyerovclaude
andcommitted
perf(gfql): fuse polars seeded RETURN-dst helper into one collect_all (#1755)
The eager form paid ~19 fixed-cost collects per execution (~70% of the fast path's time, same profile signature as the two-star OLAP lane). The same seed-first reduction now composes as one lazy plan over semi-joins on drop_nulls()'d id sets and collects once via pl.collect_all([dstn, edges]). Value-identical rows (row order may differ — documented contract); the empty-seed early return is subsumed (empty seed flows through to empty, schema-identical frames). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Y6dQEcjdazEnzuvuwf73ZL
1 parent c76af2c commit 607f42c

1 file changed

Lines changed: 26 additions & 19 deletions

File tree

graphistry/compute/chain_fast_paths.py

Lines changed: 26 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
"""
88
# ruff: noqa: E501
99

10-
from typing import Any, Dict, Optional, Tuple
10+
from typing import Any, Dict, Optional, Tuple, cast
1111

1212
from graphistry.Plottable import Plottable
1313
from .ast import ASTNode, ASTEdge, Direction
@@ -174,25 +174,32 @@ def _seeded_typed_return_dst_polars(
174174
return None
175175
from_col, to_col = (src, dst) if direction == "forward" else (dst, src)
176176

177-
# from-side seed: reduce the node frame to the seed rows, take their ids.
178-
# Membership sets are drop_nulls()'d (null ids/endpoints never link, matching
179-
# the full pipeline's joins) and passed via .implode() (Series-arg is_in is
180-
# deprecated in polars 1.42, see polars#22149).
181-
seed_nodes = nodes_df
177+
# ONE fused lazy plan, collected once (#1755 lane-3): the eager form paid a fixed
178+
# collect cost per op (~19 collects/exec = ~70% of the fast path's time). Same
179+
# reduction — seed filter -> seed out-edges -> typed-edge filter -> destination
180+
# nodes -> dangling drop/dedup — expressed with semi-joins over drop_nulls()'d id
181+
# sets (null ids/endpoints never link, matching the full pipeline's joins; the
182+
# old empty-seed early return is subsumed: an empty seed flows through to empty,
183+
# schema-identical frames). Value-identical rows; row order may differ (documented
184+
# contract; cypher without ORDER BY is unordered).
185+
lf_nodes = nodes_df.lazy()
186+
lf_edges = edges_df.lazy()
187+
seed_lf = lf_nodes
182188
for k, v in n0f.items():
183-
seed_nodes = seed_nodes.filter(pl.col(k) == v)
184-
from_ids = seed_nodes.get_column(node).drop_nulls()
185-
if from_ids.len() == 0:
186-
return nodes_df.clear(), edges_df.clear()
187-
edges = edges_df.filter(pl.col(from_col).is_in(from_ids.implode()))
189+
seed_lf = seed_lf.filter(pl.col(k) == v)
190+
from_ids_lf = seed_lf.select(pl.col(node).drop_nulls()).unique()
191+
edges_lf = lf_edges.join(from_ids_lf, left_on=from_col, right_on=node, how="semi")
188192
for k, v in ef.items(): # typed edge on the reduced frontier
189-
edges = edges.filter(pl.col(k) == v)
190-
dst_ids = edges.get_column(to_col).drop_nulls().unique()
191-
dstn = nodes_df.filter(pl.col(node).is_in(dst_ids.implode()))
193+
edges_lf = edges_lf.filter(pl.col(k) == v)
194+
dst_ids_lf = edges_lf.select(pl.col(to_col).drop_nulls().alias(node)).unique()
195+
dstn_lf = lf_nodes.join(dst_ids_lf, on=node, how="semi")
192196
for k, v in n2f.items(): # destination-node filter
193-
dstn = dstn.filter(pl.col(k) == v)
197+
dstn_lf = dstn_lf.filter(pl.col(k) == v)
194198
# drop dangling edges + dedup destination nodes (mirror the pandas tail)
195-
keep_ids = dstn.get_column(node).drop_nulls()
196-
edges = edges.filter(pl.col(to_col).is_in(keep_ids.implode()))
197-
dstn = dstn.filter(pl.col(node).is_in(edges.get_column(to_col).implode())).unique(subset=[node], maintain_order=True)
198-
return dstn, edges
199+
keep_ids_lf = dstn_lf.select(pl.col(node).drop_nulls()).unique()
200+
edges_lf = edges_lf.join(keep_ids_lf, left_on=to_col, right_on=node, how="semi")
201+
dstn_lf = dstn_lf.join(
202+
edges_lf.select(pl.col(to_col).alias(node)).unique(), on=node, how="semi"
203+
).unique(subset=[node], maintain_order=True)
204+
dstn, edges = pl.collect_all([dstn_lf, edges_lf])
205+
return cast(DataFrameT, dstn), cast(DataFrameT, edges)

0 commit comments

Comments
 (0)