|
7 | 7 | """ |
8 | 8 | # ruff: noqa: E501 |
9 | 9 |
|
10 | | -from typing import Any, Dict, Optional, Tuple |
| 10 | +from typing import Any, Dict, Optional, Tuple, cast |
11 | 11 |
|
12 | 12 | from graphistry.Plottable import Plottable |
13 | 13 | from .ast import ASTNode, ASTEdge, Direction |
@@ -174,25 +174,32 @@ def _seeded_typed_return_dst_polars( |
174 | 174 | return None |
175 | 175 | from_col, to_col = (src, dst) if direction == "forward" else (dst, src) |
176 | 176 |
|
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 |
182 | 188 | 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") |
188 | 192 | 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") |
192 | 196 | 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) |
194 | 198 | # 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