From 426179a1eedd451c14a1c6fdfc5ec0cb28b72d42 Mon Sep 17 00:00:00 2001 From: Leo Meyerovich Date: Thu, 16 Jul 2026 05:43:06 -0700 Subject: [PATCH] feat(gfql): execute connected join fast paths --- graphistry/compute/dataframe/join.py | 17 +- graphistry/compute/gfql_unified.py | 998 +++++++++++++++++- .../compute/gfql/cypher/test_lowering.py | 304 +++++- 3 files changed, 1311 insertions(+), 8 deletions(-) diff --git a/graphistry/compute/dataframe/join.py b/graphistry/compute/dataframe/join.py index 6f7081fa94..dc8fc97aa2 100644 --- a/graphistry/compute/dataframe/join.py +++ b/graphistry/compute/dataframe/join.py @@ -3,7 +3,7 @@ import operator from typing import Any, Dict, List, Optional, Sequence, Tuple, cast -from graphistry.Engine import Engine +from graphistry.Engine import Engine, POLARS_ENGINES from graphistry.compute.typing import DataFrameT, DomainT @@ -20,9 +20,15 @@ def joined_hidden_scalar_columns(frame: DataFrameT) -> DataFrameT: if suffix.startswith("__cypher_reentry_") or suffix.startswith("__gfql_hidden_"): hidden_suffixes.setdefault(suffix, []).append(column) out = frame + is_polars = "polars" in type(frame).__module__ for suffix, columns in hidden_suffixes.items(): if suffix in out.columns: continue + if is_polars: + import polars as pl + expr = pl.coalesce([pl.col(column) for column in columns]).alias(suffix) + out = out.with_columns(expr) + continue series = out[columns[0]] for column in columns[1:]: if hasattr(series, "combine_first"): @@ -44,6 +50,13 @@ def joined_alias_columns(frame: DataFrameT) -> DataFrameT: elif suffix == "id" and alias not in alias_candidates: alias_candidates[alias] = column out = frame + is_polars = "polars" in type(frame).__module__ + if is_polars and alias_candidates: + import polars as pl + return cast(DataFrameT, out.with_columns([ + pl.col(source_column).alias(alias) + for alias, source_column in alias_candidates.items() + ])) for alias, source_column in alias_candidates.items(): out = out.assign(**{alias: out[source_column]}) return out @@ -65,6 +78,8 @@ def connected_inner_join_rows( join_cols_list = list(join_cols) keep_cols_list = list(keep_cols) rhs = cast(DataFrameT, pattern_rows[keep_cols_list]) + if engine in POLARS_ENGINES: + return cast(DataFrameT, joined_rows.join(rhs, on=join_cols_list, how="inner")) if engine != Engine.CUDF: return cast(DataFrameT, joined_rows.merge(rhs, on=join_cols_list, how="inner")) diff --git a/graphistry/compute/gfql_unified.py b/graphistry/compute/gfql_unified.py index 0583073072..602e57221d 100644 --- a/graphistry/compute/gfql_unified.py +++ b/graphistry/compute/gfql_unified.py @@ -2,6 +2,7 @@ # ruff: noqa: E501 from dataclasses import replace +import pandas as pd from types import MappingProxyType from typing import Any, Dict, List, Literal, Mapping, Optional, Sequence, Tuple, Union, cast from graphistry.Plottable import Plottable @@ -63,6 +64,7 @@ joined_alias_columns as _joined_alias_columns, joined_hidden_scalar_columns as _joined_hidden_scalar_columns, ) +from graphistry.compute.filter_by_dict import filter_by_dict from graphistry.compute.gfql.ir.compilation import PhysicalPlan, PlanContext from graphistry.compute.gfql.ir.logical_plan import LogicalPlan from graphistry.compute.gfql.physical_planner import PhysicalPlanner @@ -470,6 +472,953 @@ def _optional_arm_start_nodes( return _chain_dispatch(joined_plottable, plan.post_join_chain, engine, policy, context) + +def _is_connected_fast_single_hop(edge_op: ASTEdge) -> bool: + return ( + getattr(edge_op, "direction", None) == "forward" + and getattr(edge_op, "hops", None) in (None, 1) + and getattr(edge_op, "min_hops", None) is None + and getattr(edge_op, "max_hops", None) is None + and getattr(edge_op, "output_min_hops", None) is None + and getattr(edge_op, "output_max_hops", None) is None + and getattr(edge_op, "label_node_hops", None) is None + and getattr(edge_op, "label_edge_hops", None) is None + and not bool(getattr(edge_op, "label_seeds", False)) + and not bool(getattr(edge_op, "to_fixed_point", False)) + and getattr(edge_op, "source_node_match", None) is None + and getattr(edge_op, "destination_node_match", None) is None + and getattr(edge_op, "source_node_query", None) is None + and getattr(edge_op, "destination_node_query", None) is None + and getattr(edge_op, "edge_query", None) is None + and getattr(edge_op, "_name", None) is None + and not bool(getattr(edge_op, "include_zero_hop_seed", False)) + ) + + + + +_CACHE_MISSING = object() + + +def _connected_join_filter_value_cache_key(value: Any) -> Optional[Tuple[str, str]]: + if isinstance(value, bool): + return "bool", "true" if value else "false" + if isinstance(value, str): + return "str", value + if isinstance(value, int) and not isinstance(value, bool): + return "int", str(value) + if isinstance(value, float): + return "float", repr(value) + if value is None: + return "none", "" + + predicate_value = getattr(value, "val", _CACHE_MISSING) + if predicate_value is not _CACHE_MISSING: + scalar_key = _connected_join_filter_value_cache_key(predicate_value) + if scalar_key is None: + return None + return f"{type(value).__module__}.{type(value).__qualname__}", f"{scalar_key[0]}:{scalar_key[1]}" + + regex_pattern = getattr(value, "pat", _CACHE_MISSING) + regex_case = getattr(value, "case", _CACHE_MISSING) + regex_flags = getattr(value, "flags", _CACHE_MISSING) + regex_na = getattr(value, "na", _CACHE_MISSING) + if ( + regex_pattern is not _CACHE_MISSING + and regex_case is not _CACHE_MISSING + and regex_flags is not _CACHE_MISSING + and regex_na is not _CACHE_MISSING + ): + pattern_key = _connected_join_filter_value_cache_key(regex_pattern) + case_key = _connected_join_filter_value_cache_key(regex_case) + flags_key = _connected_join_filter_value_cache_key(regex_flags) + na_key = _connected_join_filter_value_cache_key(regex_na) + if pattern_key is None or case_key is None or flags_key is None or na_key is None: + return None + return ( + f"{type(value).__module__}.{type(value).__qualname__}", + f"pat={pattern_key[0]}:{pattern_key[1]}|case={case_key[0]}:{case_key[1]}|flags={flags_key[0]}:{flags_key[1]}|na={na_key[0]}:{na_key[1]}", + ) + + predicates = getattr(value, "predicates", _CACHE_MISSING) + if predicates is not _CACHE_MISSING and isinstance(predicates, (list, tuple)): + child_keys = [] + for predicate in predicates: + child_key = _connected_join_filter_value_cache_key(predicate) + if child_key is None: + return None + child_keys.append(f"{child_key[0]}:{child_key[1]}") + return f"{type(value).__module__}.{type(value).__qualname__}", "|".join(child_keys) + + return None + + +def _connected_join_simple_filter_cache_key(filter_dict: Optional[dict]) -> Optional[Tuple[Tuple[str, str, str], ...]]: + if not filter_dict: + return () + items: List[Tuple[str, str, str]] = [] + for key, value in filter_dict.items(): + if not isinstance(key, str): + return None + value_key = _connected_join_filter_value_cache_key(value) + if value_key is None: + return None + items.append((key, value_key[0], value_key[1])) + return tuple(sorted(items)) + + +def _connected_join_cached_node_filter( + base_graph: Plottable, + nodes_obj: DataFrameT, + node_match: Optional[dict], + *, + engine: Engine, +) -> DataFrameT: + cache_key = _connected_join_simple_filter_cache_key(node_match) + if cache_key is None: + nodes = df_to_engine(nodes_obj, engine) + if engine in POLARS_ENGINES: + from graphistry.compute.gfql.lazy.engine.polars.predicates import filter_by_dict_polars + return cast(DataFrameT, filter_by_dict_polars(nodes, node_match)) + return filter_by_dict(nodes, node_match, engine=EngineAbstract(engine.value)) + + cache_attr = "_gfql_connected_join_node_filter_cache" + cache = getattr(base_graph, cache_attr, None) + if not isinstance(cache, dict): + cache = {} + try: + setattr(base_graph, cache_attr, cache) + except Exception: + cache = None + full_key = (id(nodes_obj), engine.value, cache_key) + if cache is not None and full_key in cache: + return cast(DataFrameT, cache[full_key]) + + nodes = df_to_engine(nodes_obj, engine) + if engine in POLARS_ENGINES: + from graphistry.compute.gfql.lazy.engine.polars.predicates import filter_by_dict_polars + filtered = cast(DataFrameT, filter_by_dict_polars(nodes, node_match)) + else: + filtered = filter_by_dict(nodes, node_match, engine=EngineAbstract(engine.value)) + if cache is not None: + cache[full_key] = filtered + return cast(DataFrameT, filtered) + + +def _connected_join_cached_node_ids( + base_graph: Plottable, + nodes_obj: DataFrameT, + node_matches: Sequence[Optional[dict]], + *, + node_col: str, + engine: Engine, +) -> Optional[DataFrameT]: + if engine not in POLARS_ENGINES: + return None + + match_keys: List[Tuple[Tuple[str, str, str], ...]] = [] + for node_match in node_matches: + match_key = _connected_join_simple_filter_cache_key(node_match) + if match_key is None: + return None + match_keys.append(match_key) + + cache_attr = "_gfql_connected_join_node_ids_cache" + cache = getattr(base_graph, cache_attr, None) + if not isinstance(cache, dict): + cache = {} + try: + setattr(base_graph, cache_attr, cache) + except Exception: + cache = None + full_key = (id(nodes_obj), engine.value, node_col, tuple(match_keys)) + if cache is not None and full_key in cache: + return cast(DataFrameT, cache[full_key]) + + from graphistry.compute.gfql.lazy.engine.polars.predicates import filter_by_dict_polars + + nodes = df_to_engine(nodes_obj, engine) + filtered = nodes + for node_match in node_matches: + filtered = cast(DataFrameT, filter_by_dict_polars(filtered, node_match)) + ids = cast(DataFrameT, filtered.select(node_col).unique()) + if cache is not None: + cache[full_key] = ids + return ids + + +def _connected_join_cached_edge_filter( + base_graph: Plottable, + edges_obj: DataFrameT, + edge_match: Optional[dict], + *, + engine: Engine, +) -> DataFrameT: + cache_key = _connected_join_simple_filter_cache_key(edge_match) + if cache_key is None: + edges = df_to_engine(edges_obj, engine) + if engine in POLARS_ENGINES: + from graphistry.compute.gfql.lazy.engine.polars.predicates import filter_by_dict_polars + return cast(DataFrameT, filter_by_dict_polars(edges, edge_match)) + return filter_by_dict(edges, edge_match, engine=EngineAbstract(engine.value)) + + cache_attr = "_gfql_connected_join_edge_filter_cache" + cache = getattr(base_graph, cache_attr, None) + if not isinstance(cache, dict): + cache = {} + try: + setattr(base_graph, cache_attr, cache) + except Exception: + cache = None + full_key = (id(edges_obj), engine.value, cache_key) + if cache is not None and full_key in cache: + return cast(DataFrameT, cache[full_key]) + + edges = df_to_engine(edges_obj, engine) + if engine in POLARS_ENGINES: + from graphistry.compute.gfql.lazy.engine.polars.predicates import filter_by_dict_polars + filtered = cast(DataFrameT, filter_by_dict_polars(edges, edge_match)) + else: + filtered = filter_by_dict(edges, edge_match, engine=EngineAbstract(engine.value)) + if cache is not None: + cache[full_key] = filtered + return cast(DataFrameT, filtered) + + +def _connected_join_cached_singleton_dst_source_counts( + base_graph: Plottable, + edges_obj: DataFrameT, + edge_domain: DataFrameT, + edge_match: Optional[dict], + singleton_dst: Any, + *, + src_col: str, + dst_col: str, + engine: Engine, +) -> Optional[DataFrameT]: + edge_key = _connected_join_simple_filter_cache_key(edge_match) + dst_key = _connected_join_filter_value_cache_key(singleton_dst) + if edge_key is None or dst_key is None: + return None + + cache_attr = "_gfql_connected_join_singleton_dst_source_counts_cache" + cache = getattr(base_graph, cache_attr, None) + if not isinstance(cache, dict): + cache = {} + try: + setattr(base_graph, cache_attr, cache) + except Exception: + cache = None + full_key = (id(edges_obj), engine.value, src_col, dst_col, edge_key, dst_key) + if cache is not None and full_key in cache: + return cast(DataFrameT, cache[full_key]) + + if engine in POLARS_ENGINES: + import polars as pl + counts = cast( + DataFrameT, + edge_domain + .filter(pl.col(dst_col) == singleton_dst) + .group_by(src_col) + .len("__left_count__"), + ) + else: + filtered = edge_domain[edge_domain[dst_col] == singleton_dst] + counts = cast(DataFrameT, filtered.groupby(src_col, sort=False).size().reset_index(name="__left_count__")) + + if cache is not None: + cache[full_key] = counts + return counts + + +def _connected_join_cached_first_arm_shared_counts( + base_graph: Plottable, + nodes_obj: DataFrameT, + edges_obj: DataFrameT, + first_edges: DataFrameT, + shared_ids: DataFrameT, + first_edge_match: Optional[dict], + first_start_match: Optional[dict], + second_start_match: Optional[dict], + singleton_dst: Any, + *, + shared_alias: str, + src_col: str, + dst_col: str, + engine: Engine, +) -> Optional[DataFrameT]: + if engine not in POLARS_ENGINES: + return None + + edge_key = _connected_join_simple_filter_cache_key(first_edge_match) + first_start_key = _connected_join_simple_filter_cache_key(first_start_match) + second_start_key = _connected_join_simple_filter_cache_key(second_start_match) + dst_key = _connected_join_filter_value_cache_key(singleton_dst) + if edge_key is None or first_start_key is None or second_start_key is None or dst_key is None: + return None + + cache_attr = "_gfql_connected_join_first_arm_shared_counts_cache" + cache = getattr(base_graph, cache_attr, None) + if not isinstance(cache, dict): + cache = {} + try: + setattr(base_graph, cache_attr, cache) + except Exception: + cache = None + full_key = ( + id(nodes_obj), + id(edges_obj), + engine.value, + shared_alias, + src_col, + dst_col, + edge_key, + first_start_key, + second_start_key, + dst_key, + ) + if cache is not None and full_key in cache: + return cast(DataFrameT, cache[full_key]) + + import polars as pl + + counts = ( + first_edges + .filter(pl.col(dst_col) == singleton_dst) + .join(shared_ids, left_on=src_col, right_on=shared_alias, how="semi") + .group_by(src_col) + .len("__left_count__") + .rename({src_col: shared_alias}) + ) + if cache is not None: + cache[full_key] = counts + return cast(DataFrameT, counts) + + +def _connected_join_cached_second_arm_group_rows( + base_graph: Plottable, + nodes_obj: DataFrameT, + edges_obj: DataFrameT, + second_edges: DataFrameT, + shared_ids: DataFrameT, + second_leaf_ids: DataFrameT, + second_leaf_nodes: DataFrameT, + first_start_match: Optional[dict], + second_start_match: Optional[dict], + second_leaf_match: Optional[dict], + second_edge_match: Optional[dict], + second_leaf_singleton: Any, + group_prop_refs: List[Tuple[str, str]], + output_group_keys: List[str], + *, + node_col: str, + src_col: str, + dst_col: str, + engine: Engine, +) -> Optional[DataFrameT]: + if engine not in POLARS_ENGINES: + return None + + first_start_key = _connected_join_simple_filter_cache_key(first_start_match) + second_start_key = _connected_join_simple_filter_cache_key(second_start_match) + second_leaf_key = _connected_join_simple_filter_cache_key(second_leaf_match) + second_edge_key = _connected_join_simple_filter_cache_key(second_edge_match) + if first_start_key is None or second_start_key is None or second_leaf_key is None or second_edge_key is None: + return None + + prop_key = tuple((str(out_col), str(prop)) for out_col, prop in group_prop_refs) + output_key = tuple(str(key) for key in output_group_keys) + cache_attr = "_gfql_connected_join_second_arm_group_rows_cache" + cache = getattr(base_graph, cache_attr, None) + if not isinstance(cache, dict): + cache = {} + try: + setattr(base_graph, cache_attr, cache) + except Exception: + cache = None + full_key = ( + id(nodes_obj), + id(edges_obj), + engine.value, + node_col, + src_col, + dst_col, + first_start_key, + second_start_key, + second_leaf_key, + second_edge_key, + prop_key, + output_key, + ) + if cache is not None and full_key in cache: + return cast(DataFrameT, cache[full_key]) + + import polars as pl + + if second_leaf_singleton is _CACHE_MISSING: + right_base = ( + second_edges + .join(shared_ids, left_on=src_col, right_on=node_col, how="semi") + .join(second_leaf_ids, left_on=dst_col, right_on=node_col, how="semi") + ) + else: + right_base = ( + second_edges + .filter(pl.col(dst_col) == second_leaf_singleton) + .join(shared_ids, left_on=src_col, right_on=node_col, how="semi") + ) + if group_prop_refs: + lookup_key = "__gfql_fast_second_leaf_id__" + lookup_exprs = [pl.col(node_col).alias(lookup_key)] + [pl.col(prop).alias(out_col) for out_col, prop in group_prop_refs] + second_lookup = second_leaf_nodes.select(lookup_exprs) + right_base = right_base.join(second_lookup, left_on=dst_col, right_on=lookup_key, how="inner") + right_rows = right_base.select([pl.col(src_col)] + [pl.col(key) for key in output_group_keys]) + if cache is not None: + cache[full_key] = right_rows + return cast(DataFrameT, right_rows) + + +def _property_ref(expr: Any, valid_aliases: Sequence[str]) -> Optional[Tuple[str, str]]: + if not isinstance(expr, str) or "." not in expr: + return None + alias, prop = expr.split(".", 1) + if alias not in valid_aliases or not prop: + return None + return alias, prop + + +def _connected_join_post_property_columns(plan: ConnectedMatchJoinPlan, alias: str) -> List[str]: + prefix = f"{alias}." + out: List[str] = [] + + def walk(value: Any) -> None: + if isinstance(value, str): + if value.startswith(prefix): + prop = value[len(prefix):] + if prop and prop not in out: + out.append(prop) + return + if isinstance(value, Mapping): + for item in value.values(): + walk(item) + return + if isinstance(value, (list, tuple)): + for item in value: + walk(item) + + for op in plan.post_join_chain.chain: + if isinstance(op, ASTCall): + walk(op.params) + return out + + + +def _connected_join_two_star_fast_grouped_count( + base_graph: Plottable, + plan: ConnectedMatchJoinPlan, + *, + engine: Engine, +) -> Optional[DataFrameT]: + if len(plan.pattern_chains) != 2 or len(plan.pattern_shared_node_aliases) != 1: + return None + if len(plan.pattern_shared_node_aliases[0]) != 1: + return None + shared_alias = plan.pattern_shared_node_aliases[0][0] + + parsed = [] + for pattern_chain in plan.pattern_chains: + if pattern_chain.where: + return None + ops = list(pattern_chain.chain) + if len(ops) != 3: + return None + start_op, edge_op, end_op = ops + if not isinstance(start_op, ASTNode) or not isinstance(edge_op, ASTEdge) or not isinstance(end_op, ASTNode): + return None + if getattr(start_op, "query", None) is not None or getattr(end_op, "query", None) is not None: + return None + if not _is_connected_fast_single_hop(edge_op): + return None + start_alias = getattr(start_op, "_name", None) + end_alias = getattr(end_op, "_name", None) + if start_alias != shared_alias or not isinstance(end_alias, str): + return None + parsed.append((start_op, edge_op, end_op, end_alias)) + + first_start, first_edge, first_end, first_end_alias = parsed[0] + second_start, second_edge, second_end, second_end_alias = parsed[1] + if first_end_alias == second_end_alias: + return None + + post_ops = [op for op in plan.post_join_chain.chain if isinstance(op, ASTCall)] + if len(post_ops) not in (2, 3, 4): + return None + if post_ops[0].function != "with_" or post_ops[1].function != "group_by": + return None + suffix = post_ops[2:] + order_call: Optional[ASTCall] = None + limit_call: Optional[ASTCall] = None + select_call: Optional[ASTCall] = None + for op in suffix: + if op.function == "order_by" and order_call is None and limit_call is None and select_call is None: + order_call = op + elif op.function == "limit" and limit_call is None and select_call is None: + limit_call = op + elif op.function == "select" and select_call is None: + select_call = op + else: + return None + + with_items_raw = post_ops[0].params.get("items") + group_keys_raw = post_ops[1].params.get("keys") + aggs_raw = post_ops[1].params.get("aggregations") + if not isinstance(with_items_raw, list) or not isinstance(group_keys_raw, list) or not isinstance(aggs_raw, list): + return None + if len(aggs_raw) != 1: + return None + agg = aggs_raw[0] + if not isinstance(agg, (tuple, list)) or len(agg) != 3 or str(agg[1]).lower() != "count": + return None + agg_alias = str(agg[0]) + agg_input = agg[2] + + with_items: Dict[str, Any] = {} + for item in with_items_raw: + if not isinstance(item, (tuple, list)) or len(item) != 2 or not isinstance(item[0], str): + return None + with_items[item[0]] = item[1] + if not isinstance(agg_input, str) or with_items.get(agg_input) != shared_alias: + return None + + group_keys = [str(key) for key in group_keys_raw] + group_prop_refs: List[Tuple[str, str]] = [] + output_group_keys: List[str] = [] + for key in group_keys: + expr = with_items.get(key) + if key == "__cypher_group__" and expr == 1: + continue + prop_ref = _property_ref(expr, (second_end_alias,)) + if prop_ref is None: + return None + output_group_keys.append(key) + group_prop_refs.append((key, prop_ref[1])) + + order_keys: List[Tuple[str, bool]] = [] + if order_call is not None: + raw_order = order_call.params.get("keys") + if not isinstance(raw_order, list): + return None + available = set(output_group_keys) | {agg_alias} + for item in raw_order: + if not isinstance(item, (tuple, list)) or len(item) != 2 or not isinstance(item[0], str): + return None + if item[0] not in available: + return None + direction = str(item[1]).lower() + if direction not in {"asc", "ascending", "desc", "descending"}: + return None + order_keys.append((item[0], direction in {"desc", "descending"})) + + limit_value: Optional[int] = None + if limit_call is not None: + raw_limit = limit_call.params.get("value") + if not isinstance(raw_limit, int) or raw_limit < 0: + return None + limit_value = raw_limit + + select_items: Optional[List[Tuple[str, str]]] = None + if select_call is not None: + raw_select = select_call.params.get("items") + if not isinstance(raw_select, list): + return None + select_items = [] + available = set(output_group_keys) | {agg_alias} + for item in raw_select: + if not isinstance(item, (tuple, list)) or len(item) != 2 or not isinstance(item[0], str) or not isinstance(item[1], str): + return None + if item[0] not in available: + return None + select_items.append((item[0], item[1])) + + nodes_obj = getattr(base_graph, "_nodes", None) + edges_obj = getattr(base_graph, "_edges", None) + node_col = getattr(base_graph, "_node", None) + src_col = getattr(base_graph, "_source", None) + dst_col = getattr(base_graph, "_destination", None) + if nodes_obj is None or edges_obj is None or node_col is None or src_col is None or dst_col is None: + return None + node_col = str(node_col) + src_col = str(src_col) + dst_col = str(dst_col) + if node_col not in nodes_obj.columns or src_col not in edges_obj.columns or dst_col not in edges_obj.columns: + return None + + nodes_source = cast(DataFrameT, nodes_obj) + nodes = df_to_engine(nodes_source, engine) + edges = cast(DataFrameT, edges_obj) + + if engine in POLARS_ENGINES: + import polars as pl + from graphistry.compute.gfql.lazy.engine.polars.predicates import filter_by_dict_polars + + if isinstance(nodes_obj, (pl.DataFrame, pl.LazyFrame)): + shared_ids = _connected_join_cached_node_ids( + base_graph, + nodes_source, + [cast(Optional[dict], first_start.filter_dict), cast(Optional[dict], second_start.filter_dict)], + node_col=node_col, + engine=engine, + ) + first_leaf_ids = _connected_join_cached_node_ids( + base_graph, + nodes_source, + [cast(Optional[dict], first_end.filter_dict)], + node_col=node_col, + engine=engine, + ) + second_leaf_ids = _connected_join_cached_node_ids( + base_graph, + nodes_source, + [cast(Optional[dict], second_end.filter_dict)], + node_col=node_col, + engine=engine, + ) + first_leaf_nodes = _connected_join_cached_node_filter(base_graph, nodes_source, cast(Optional[dict], first_end.filter_dict), engine=engine) + second_leaf_nodes = _connected_join_cached_node_filter(base_graph, nodes_source, cast(Optional[dict], second_end.filter_dict), engine=engine) + if shared_ids is None or first_leaf_ids is None or second_leaf_ids is None: + return None + else: + shared_nodes = filter_by_dict_polars(nodes, cast(Optional[dict], first_start.filter_dict)) + shared_nodes = filter_by_dict_polars(shared_nodes, cast(Optional[dict], second_start.filter_dict)) + first_leaf_nodes = filter_by_dict_polars(nodes, cast(Optional[dict], first_end.filter_dict)) + second_leaf_nodes = filter_by_dict_polars(nodes, cast(Optional[dict], second_end.filter_dict)) + shared_ids = shared_nodes.select(node_col).unique() + first_leaf_ids = first_leaf_nodes.select(node_col).unique() + second_leaf_ids = second_leaf_nodes.select(node_col).unique() + first_edges = _connected_join_cached_edge_filter(base_graph, edges, cast(Optional[dict], first_edge.edge_match), engine=engine) + second_edges = _connected_join_cached_edge_filter(base_graph, edges, cast(Optional[dict], second_edge.edge_match), engine=engine) + + for _, prop in group_prop_refs: + if prop not in second_leaf_nodes.columns: + return None + + def singleton_id(ids: Any) -> Any: + if not isinstance(ids, pl.DataFrame) or len(ids) != 1: + return _CACHE_MISSING + return ids.get_column(node_col)[0] + + first_leaf_singleton = singleton_id(first_leaf_ids) + second_leaf_singleton = singleton_id(second_leaf_ids) + + cached_left_counts: Optional[DataFrameT] = None + if first_leaf_singleton is not _CACHE_MISSING: + shared_ids_for_left = shared_ids.rename({node_col: shared_alias}) if node_col != shared_alias else shared_ids + cached_left_counts = _connected_join_cached_first_arm_shared_counts( + base_graph, + nodes_source, + edges, + first_edges, + shared_ids_for_left, + cast(Optional[dict], first_edge.edge_match), + cast(Optional[dict], first_start.filter_dict), + cast(Optional[dict], second_start.filter_dict), + first_leaf_singleton, + shared_alias=shared_alias, + src_col=src_col, + dst_col=dst_col, + engine=engine, + ) + if cached_left_counts is None: + cached_left_counts = _connected_join_cached_singleton_dst_source_counts( + base_graph, + edges, + first_edges, + cast(Optional[dict], first_edge.edge_match), + first_leaf_singleton, + src_col=src_col, + dst_col=dst_col, + engine=engine, + ) + if cached_left_counts is not None: + if shared_alias in cached_left_counts.columns: + left_counts = cached_left_counts + else: + left_counts = ( + cached_left_counts + .join(shared_ids, left_on=src_col, right_on=node_col, how="semi") + .rename({src_col: shared_alias}) + ) + elif first_leaf_singleton is _CACHE_MISSING: + left_edges = ( + first_edges + .join(shared_ids, left_on=src_col, right_on=node_col, how="semi") + .join(first_leaf_ids, left_on=dst_col, right_on=node_col, how="semi") + ) + left_counts = ( + left_edges + .group_by(src_col) + .len("__left_count__") + .rename({src_col: shared_alias}) + ) + else: + left_edges = ( + first_edges + .filter(pl.col(dst_col) == first_leaf_singleton) + .join(shared_ids, left_on=src_col, right_on=node_col, how="semi") + ) + left_counts = ( + left_edges + .group_by(src_col) + .len("__left_count__") + .rename({src_col: shared_alias}) + ) + cached_right_rows = _connected_join_cached_second_arm_group_rows( + base_graph, + nodes_source, + edges, + second_edges, + shared_ids, + second_leaf_ids, + second_leaf_nodes, + cast(Optional[dict], first_start.filter_dict), + cast(Optional[dict], second_start.filter_dict), + cast(Optional[dict], second_end.filter_dict), + cast(Optional[dict], second_edge.edge_match), + second_leaf_singleton, + group_prop_refs, + output_group_keys, + node_col=node_col, + src_col=src_col, + dst_col=dst_col, + engine=engine, + ) + if cached_right_rows is not None: + right_rows = cached_right_rows if src_col == shared_alias else cached_right_rows.rename({src_col: shared_alias}) + else: + if second_leaf_singleton is _CACHE_MISSING: + right_base = ( + second_edges + .join(shared_ids, left_on=src_col, right_on=node_col, how="semi") + .join(second_leaf_ids, left_on=dst_col, right_on=node_col, how="semi") + ) + else: + right_base = ( + second_edges + .filter(pl.col(dst_col) == second_leaf_singleton) + .join(shared_ids, left_on=src_col, right_on=node_col, how="semi") + ) + if group_prop_refs: + lookup_key = "__gfql_fast_second_leaf_id__" + lookup_exprs = [pl.col(node_col).alias(lookup_key)] + [pl.col(prop).alias(out_col) for out_col, prop in group_prop_refs] + second_lookup = second_leaf_nodes.select(lookup_exprs) + right_base = right_base.join(second_lookup, left_on=dst_col, right_on=lookup_key, how="inner") + right_rows = right_base.select([pl.col(src_col).alias(shared_alias)] + [pl.col(key) for key in output_group_keys]) + if not output_group_keys and len(left_counts) > 0 and bool(left_counts.select((pl.col("__left_count__") == 1).all()).item()): + matched_rows = right_rows.join(left_counts.select(shared_alias), on=shared_alias, how="semi") + out_df = matched_rows.select(pl.len().cast(pl.Int64).alias(agg_alias)) + else: + joined = right_rows.join(left_counts, on=shared_alias, how="inner") + if len(joined) == 0: + return cast(DataFrameT, joined.select([])) + if output_group_keys: + out_df = joined.group_by(output_group_keys, maintain_order=True).agg(pl.col("__left_count__").sum().cast(pl.Int64).alias(agg_alias)) + else: + out_df = joined.select(pl.col("__left_count__").sum().cast(pl.Int64).alias(agg_alias)) + if order_keys: + out_df = out_df.sort([key for key, _ in order_keys], descending=[desc for _, desc in order_keys]) + if limit_value is not None: + out_df = out_df.head(limit_value) + if select_items is not None: + out_df = out_df.select([pl.col(src).alias(dst) for src, dst in select_items]) + return cast(DataFrameT, out_df) + + filter_engine = EngineAbstract(engine.value) + shared_nodes = filter_by_dict(nodes, cast(Optional[dict], first_start.filter_dict), engine=filter_engine) + shared_nodes = filter_by_dict(shared_nodes, cast(Optional[dict], second_start.filter_dict), engine=filter_engine) + first_leaf_nodes = filter_by_dict(nodes, cast(Optional[dict], first_end.filter_dict), engine=filter_engine) + second_leaf_nodes = filter_by_dict(nodes, cast(Optional[dict], second_end.filter_dict), engine=filter_engine) + first_edges = _connected_join_cached_edge_filter(base_graph, edges, cast(Optional[dict], first_edge.edge_match), engine=engine) + second_edges = _connected_join_cached_edge_filter(base_graph, edges, cast(Optional[dict], second_edge.edge_match), engine=engine) + + shared_ids = shared_nodes[node_col].drop_duplicates() + first_leaf_ids = first_leaf_nodes[node_col].drop_duplicates() + second_leaf_ids = second_leaf_nodes[node_col].drop_duplicates() + for _, prop in group_prop_refs: + if prop not in second_leaf_nodes.columns: + return None + + left_edges = first_edges[first_edges[src_col].isin(shared_ids) & first_edges[dst_col].isin(first_leaf_ids)] + left_counts = left_edges.groupby(src_col, sort=False).size().reset_index(name="__left_count__").rename(columns={src_col: shared_alias}) + right_base = second_edges[second_edges[src_col].isin(shared_ids) & second_edges[dst_col].isin(second_leaf_ids)] + if group_prop_refs: + lookup_key = "__gfql_fast_second_leaf_id__" + prop_cols = [] + for _, prop in group_prop_refs: + if prop not in prop_cols: + prop_cols.append(prop) + second_lookup = second_leaf_nodes[[node_col] + prop_cols].drop_duplicates(subset=[node_col]).rename(columns={node_col: lookup_key}) + for out_col, prop in group_prop_refs: + second_lookup[out_col] = second_lookup[prop] + right_base = right_base.merge(second_lookup[[lookup_key] + output_group_keys], left_on=dst_col, right_on=lookup_key, how="inner") + right_rows = right_base[[src_col] + output_group_keys].rename(columns={src_col: shared_alias}) + joined = right_rows.merge(left_counts, on=shared_alias, how="inner") + if len(joined) == 0: + return df_cons(engine)() + if output_group_keys: + out_df = joined.groupby(output_group_keys, sort=False, dropna=False)["__left_count__"].sum().reset_index(name=agg_alias) + else: + out_df = pd.DataFrame({agg_alias: [int(joined["__left_count__"].sum())]}) + if order_keys: + out_df = cast(DataFrameT, out_df.sort_values(by=[key for key, _ in order_keys], ascending=[not desc for _, desc in order_keys])) + if limit_value is not None: + out_df = cast(DataFrameT, out_df.head(limit_value)) + if select_items is not None: + out_df = out_df[[src for src, _ in select_items]].rename(columns={src: dst for src, dst in select_items}) + return df_to_engine(out_df.reset_index(drop=True), engine) + + +def _connected_join_two_star_fast_rows( + base_graph: Plottable, + plan: ConnectedMatchJoinPlan, + *, + engine: Engine, +) -> Optional[DataFrameT]: + """Fast rows for T1 two-star connected joins. + + Supported shape: ``(p)-[r1]->(i), (p)-[r2]->(c)`` where both arms are fixed + forward one-hop patterns, share the same start node alias, have no residual + per-arm row filters, and downstream projections need at most properties from + the second leaf alias. The returned frame preserves row multiplicity, then + the normal row pipeline handles RETURN/GROUP/ORDER/LIMIT. + """ + if len(plan.pattern_chains) != 2 or len(plan.pattern_shared_node_aliases) != 1: + return None + if len(plan.pattern_shared_node_aliases[0]) != 1: + return None + shared_alias = plan.pattern_shared_node_aliases[0][0] + + parsed = [] + for pattern_chain in plan.pattern_chains: + if pattern_chain.where: + return None + ops = list(pattern_chain.chain) + if len(ops) != 3: + return None + start_op, edge_op, end_op = ops + if not isinstance(start_op, ASTNode) or not isinstance(edge_op, ASTEdge) or not isinstance(end_op, ASTNode): + return None + if getattr(start_op, "query", None) is not None or getattr(end_op, "query", None) is not None: + return None + if not _is_connected_fast_single_hop(edge_op): + return None + start_alias = getattr(start_op, "_name", None) + end_alias = getattr(end_op, "_name", None) + if start_alias != shared_alias or not isinstance(end_alias, str): + return None + parsed.append((start_op, edge_op, end_op, end_alias)) + + first_start, first_edge, first_end, first_end_alias = parsed[0] + second_start, second_edge, second_end, second_end_alias = parsed[1] + if first_end_alias == second_end_alias: + return None + + attach_aliases = plan.pattern_attach_prop_aliases + if not attach_aliases or len(attach_aliases) != 2: + return None + normalized_attach_aliases: List[Tuple[str, ...]] = [] + for aliases in attach_aliases: + if aliases is None: + return None + normalized_attach_aliases.append(aliases) + needed_attach_aliases = {alias for aliases in normalized_attach_aliases for alias in aliases} + if any(alias != second_end_alias for alias in needed_attach_aliases): + return None + second_props_needed = _connected_join_post_property_columns(plan, second_end_alias) + if second_end_alias in needed_attach_aliases and not second_props_needed: + return None + + nodes_obj = getattr(base_graph, "_nodes", None) + edges_obj = getattr(base_graph, "_edges", None) + node_col = getattr(base_graph, "_node", None) + src_col = getattr(base_graph, "_source", None) + dst_col = getattr(base_graph, "_destination", None) + if nodes_obj is None or edges_obj is None or node_col is None or src_col is None or dst_col is None: + return None + node_col = str(node_col) + src_col = str(src_col) + dst_col = str(dst_col) + if node_col not in nodes_obj.columns or src_col not in edges_obj.columns or dst_col not in edges_obj.columns: + return None + + nodes = df_to_engine(cast(DataFrameT, nodes_obj), engine) + edges = cast(DataFrameT, edges_obj) + + if engine in POLARS_ENGINES: + import polars as pl + from graphistry.compute.gfql.lazy.engine.polars.predicates import filter_by_dict_polars + + shared_nodes = filter_by_dict_polars(nodes, cast(Optional[dict], first_start.filter_dict)) + shared_nodes = filter_by_dict_polars(shared_nodes, cast(Optional[dict], second_start.filter_dict)) + first_leaf_nodes = filter_by_dict_polars(nodes, cast(Optional[dict], first_end.filter_dict)) + second_leaf_nodes = filter_by_dict_polars(nodes, cast(Optional[dict], second_end.filter_dict)) + first_edges = _connected_join_cached_edge_filter(base_graph, edges, cast(Optional[dict], first_edge.edge_match), engine=engine) + second_edges = _connected_join_cached_edge_filter(base_graph, edges, cast(Optional[dict], second_edge.edge_match), engine=engine) + + shared_ids = shared_nodes.select(node_col).unique() + first_leaf_ids = first_leaf_nodes.select(node_col).unique() + second_leaf_ids = second_leaf_nodes.select(node_col).unique() + second_props = [col for col in second_props_needed if col in second_leaf_nodes.columns] + + left_rows = ( + first_edges + .join(shared_ids, left_on=src_col, right_on=node_col, how="semi") + .join(first_leaf_ids, left_on=dst_col, right_on=node_col, how="semi") + .select(pl.col(src_col).alias(shared_alias)) + ) + right_base = ( + second_edges + .join(shared_ids, left_on=src_col, right_on=node_col, how="semi") + .join(second_leaf_ids, left_on=dst_col, right_on=node_col, how="semi") + ) + if second_props: + second_lookup_key = "__gfql_fast_second_leaf_id__" + second_lookup = second_leaf_nodes.select([node_col] + second_props).rename({node_col: second_lookup_key}) + second_lookup = second_lookup.rename({col: f"{second_end_alias}.{col}" for col in second_props}) + right_base = right_base.join(second_lookup, left_on=dst_col, right_on=second_lookup_key, how="inner") + right_select = [pl.col(src_col).alias(shared_alias)] + [pl.col(f"{second_end_alias}.{col}") for col in second_props] + right_rows = right_base.select(right_select) + return cast(DataFrameT, left_rows.join(right_rows, on=shared_alias, how="inner")) + + filter_engine = EngineAbstract(engine.value) + shared_nodes = filter_by_dict(nodes, cast(Optional[dict], first_start.filter_dict), engine=filter_engine) + shared_nodes = filter_by_dict(shared_nodes, cast(Optional[dict], second_start.filter_dict), engine=filter_engine) + first_leaf_nodes = filter_by_dict(nodes, cast(Optional[dict], first_end.filter_dict), engine=filter_engine) + second_leaf_nodes = filter_by_dict(nodes, cast(Optional[dict], second_end.filter_dict), engine=filter_engine) + first_edges = _connected_join_cached_edge_filter(base_graph, edges, cast(Optional[dict], first_edge.edge_match), engine=engine) + second_edges = _connected_join_cached_edge_filter(base_graph, edges, cast(Optional[dict], second_edge.edge_match), engine=engine) + + shared_ids = shared_nodes[node_col].drop_duplicates() + first_leaf_ids = first_leaf_nodes[node_col].drop_duplicates() + second_leaf_ids = second_leaf_nodes[node_col].drop_duplicates() + second_props = [col for col in second_props_needed if col in second_leaf_nodes.columns] + + left_edges = first_edges[first_edges[src_col].isin(shared_ids) & first_edges[dst_col].isin(first_leaf_ids)] + left_rows = left_edges[[src_col]].rename(columns={src_col: shared_alias}) + right_base = second_edges[second_edges[src_col].isin(shared_ids) & second_edges[dst_col].isin(second_leaf_ids)] + if second_props: + second_lookup_key = "__gfql_fast_second_leaf_id__" + second_lookup_cols = [node_col] + second_props + second_lookup = second_leaf_nodes[second_lookup_cols].drop_duplicates(subset=[node_col]).rename( + columns={ + node_col: second_lookup_key, + **{col: f"{second_end_alias}.{col}" for col in second_props}, + } + ) + right_base = right_base.merge(second_lookup, left_on=dst_col, right_on=second_lookup_key, how="inner") + right_cols = [src_col] + [f"{second_end_alias}.{col}" for col in second_props] + right_rows = right_base[right_cols].rename(columns={src_col: shared_alias}) + return cast(DataFrameT, left_rows.merge(right_rows, on=shared_alias, how="inner")) + def _apply_connected_match_join( base_graph: Plottable, plan: ConnectedMatchJoinPlan, @@ -478,17 +1427,42 @@ def _apply_connected_match_join( policy: Optional[PolicyDict], context: ExecutionContext, ) -> Plottable: - from graphistry.compute.ast import ASTCall, serialize_binding_ops + from graphistry.compute.ast import ASTCall, ASTNode as _ASTNode, serialize_binding_ops requested_engine = resolve_engine(cast(Any, engine), base_graph) dispatch_engine: Union[EngineAbstract, str] = engine df_ctor = df_cons(requested_engine) node_col = getattr(base_graph, "_node", "id") + fast_grouped_count = _connected_join_two_star_fast_grouped_count(base_graph, plan, engine=requested_engine) + if fast_grouped_count is not None: + out = base_graph.bind() + out._nodes = fast_grouped_count + out._edges = df_ctor() + return out + + fast_rows = _connected_join_two_star_fast_rows(base_graph, plan, engine=requested_engine) + if fast_rows is not None: + if len(fast_rows) == 0: + out = base_graph.bind() + out._nodes = df_ctor() + out._edges = df_ctor() + return out + fast_rows = _joined_hidden_scalar_columns(fast_rows) + fast_rows = _joined_alias_columns(fast_rows) + joined_plottable = base_graph.bind() + joined_plottable._nodes = fast_rows + joined_plottable._edges = df_ctor() + return _chain_dispatch(joined_plottable, plan.post_join_chain, dispatch_engine, policy, context) + joined_rows: Optional[DataFrameT] = None + pattern_attach_prop_aliases = plan.pattern_attach_prop_aliases or tuple(None for _ in plan.pattern_chains) for idx, pattern_chain in enumerate(plan.pattern_chains): + rows_params: Dict[str, Any] = {"binding_ops": serialize_binding_ops(pattern_chain.chain)} + if idx < len(pattern_attach_prop_aliases) and pattern_attach_prop_aliases[idx] is not None: + rows_params["attach_prop_aliases"] = list(cast(Tuple[str, ...], pattern_attach_prop_aliases[idx])) with_rows = Chain( - list(pattern_chain.chain) + [ASTCall("rows", {"binding_ops": serialize_binding_ops(pattern_chain.chain)})], + list(pattern_chain.chain) + [ASTCall("rows", rows_params)], where=pattern_chain.where, ) pattern_result = _chain_dispatch(base_graph, with_rows, dispatch_engine, policy, context) @@ -498,16 +1472,30 @@ def _apply_connected_match_join( out._nodes = df_ctor() out._edges = df_ctor() return out - pattern_rows = cast(DataFrameT, pattern_rows[_binding_join_columns(pattern_rows)]) + node_aliases = [ + cast(str, getattr(op, "_name")) + for op in pattern_chain.chain + if isinstance(op, _ASTNode) and isinstance(getattr(op, "_name", None), str) + ] + keep_binding_columns = _binding_join_columns(pattern_rows) + [ + alias for alias in node_aliases if alias in pattern_rows.columns + ] + pattern_rows = cast(DataFrameT, pattern_rows[keep_binding_columns]) if joined_rows is None: joined_rows = pattern_rows continue shared_aliases = plan.pattern_shared_node_aliases[idx - 1] join_cols = [ - f"{alias}.{node_col}" + alias for alias in shared_aliases - if f"{alias}.{node_col}" in joined_rows.columns and f"{alias}.{node_col}" in pattern_rows.columns + if alias in joined_rows.columns and alias in pattern_rows.columns ] + if not join_cols: + join_cols = [ + f"{alias}.{node_col}" + for alias in shared_aliases + if f"{alias}.{node_col}" in joined_rows.columns and f"{alias}.{node_col}" in pattern_rows.columns + ] if not join_cols: raise GFQLValidationError( ErrorCode.E108, diff --git a/graphistry/tests/compute/gfql/cypher/test_lowering.py b/graphistry/tests/compute/gfql/cypher/test_lowering.py index 5f84f963d1..7a7057d47b 100644 --- a/graphistry/tests/compute/gfql/cypher/test_lowering.py +++ b/graphistry/tests/compute/gfql/cypher/test_lowering.py @@ -7,7 +7,10 @@ import graphistry from typing import Any, Callable, Dict, List, Optional, Tuple, cast +from graphistry.Engine import Engine + from graphistry.compute.ast import ASTCall, ASTNode, ASTEdge, ASTEdgeForward, ASTEdgeReverse, ASTEdgeUndirected +from graphistry.compute.chain import Chain from graphistry.compute.exceptions import ErrorCode, GFQLSyntaxError, GFQLTypeError, GFQLValidationError from graphistry.compute.predicates.is_in import IsIn from graphistry.compute.gfql.same_path_types import NODE_IDENTITY_COLUMN, col, compare @@ -39,7 +42,14 @@ from graphistry.plugins.networkx.policy import NETWORKX_SCIPY_EXTRA_REQUIREMENTS, NETWORKX_VERSION_SPEC, SCIPY_VERSION_SPEC from graphistry.compute.gfql.cypher.ast import ExpressionText, OrderByClause, OrderItem, ReturnClause, ReturnItem, SourceSpan from graphistry.compute.gfql.cypher.lowering import CompiledCypherExecutionExtras, CompiledCypherGraphQuery -from graphistry.compute.gfql.cypher.lowering import _logical_plan_route_for_query +from graphistry.compute.gfql.cypher.lowering import ConnectedMatchJoinPlan, _logical_plan_route_for_query +from graphistry.compute.gfql_unified import ( + _connected_join_cached_edge_filter, + _connected_join_cached_singleton_dst_source_counts, + _connected_join_post_property_columns, + _connected_join_two_star_fast_grouped_count, + _connected_join_two_star_fast_rows, +) from graphistry.compute.gfql.frontends.cypher.binder import FrontendBinder from graphistry.compute.gfql.ir.bound_ir import BoundIR, BoundQueryPart, SemanticTable from graphistry.compute.gfql.ir.compilation import PlanContext @@ -14933,7 +14943,6 @@ def test_issue_1273_multi_source_grouped_aggregate(agg: str, expected: list) -> assert result._nodes.to_dict(orient="records") == expected - def test_issue_1712_connected_comma_pattern_where_intersects() -> None: """#1712: a connected comma-pattern sharing a node alias with a WHERE on a leaf alias must intersect both patterns (the WHERE was silently dropped on the @@ -15059,6 +15068,272 @@ def test_t1_connected_comma_pushes_q7_range_filters_before_join() -> None: assert any(entry.get("c", {}).get("country") == "France" for entry in filters_by_alias) +@pytest.mark.parametrize("engine", ["pandas", "polars"]) +def test_t1_connected_comma_q5_global_count_runs_on_pandas_and_polars(engine: str) -> None: + if engine == "polars": + pytest.importorskip("polars") + query = ( + "MATCH (p {node_type:'Person'})-[{rel:'HAS_INTEREST'}]->(i {node_type:'Interest'}), " + "(p)-[{rel:'LIVES_IN'}]->(c {node_type:'City'}) " + "WHERE toLower(i.interest) = toLower('fine dining') " + "AND toLower(p.gender) = toLower('male') " + "AND c.city = 'London' AND c.country = 'United Kingdom' " + "RETURN count(p) AS numPersons" + ) + + result = _mk_graph_benchmark_t1_shape_graph().gfql(query, engine=engine) + assert _to_pandas_df(result._nodes).to_dict(orient="records") == [{"numPersons": 1}] + + +@pytest.mark.parametrize("engine", ["pandas", "polars"]) +def test_t1_connected_comma_q6_q7_grouped_count_runs_on_pandas_and_polars(engine: str) -> None: + if engine == "polars": + pytest.importorskip("polars") + query = ( + "MATCH (p {node_type:'Person'})-[{rel:'HAS_INTEREST'}]->(i {node_type:'Interest'}), " + "(p)-[{rel:'LIVES_IN'}]->(c {node_type:'City'}) " + "WHERE toLower(i.interest) = toLower('fine dining') " + "AND p.age >= 23 AND p.age <= 30 " + "AND c.country = 'United Kingdom' " + "RETURN count(p) AS numPersons, c.state AS state, c.country AS country " + "ORDER BY state" + ) + + result = _mk_graph_benchmark_t1_shape_graph().gfql(query, engine=engine) + assert _to_pandas_df(result._nodes).to_dict(orient="records") == [ + {"numPersons": 2, "state": "England", "country": "United Kingdom"} + ] + + +@pytest.mark.parametrize("engine", ["pandas", "polars"]) +def test_t1_connected_comma_two_star_fast_path_preserves_multiplicity(engine: str) -> None: + if engine == "polars": + pytest.importorskip("polars") + nodes = pd.DataFrame({ + "id": [0, 10, 20], + "node_type": ["Person", "City", "Interest"], + "gender": ["male", None, None], + "city": [None, "London", None], + "country": [None, "United Kingdom", None], + "state": [None, "England", None], + "interest": [None, None, "Fine Dining"], + }) + edges = pd.DataFrame({ + "s": [0, 0, 0, 0], + "d": [20, 20, 10, 10], + "rel": ["HAS_INTEREST", "HAS_INTEREST", "LIVES_IN", "LIVES_IN"], + }) + query = ( + "MATCH (p {node_type:'Person'})-[{rel:'HAS_INTEREST'}]->(i {node_type:'Interest'}), " + "(p)-[{rel:'LIVES_IN'}]->(c {node_type:'City'}) " + "WHERE toLower(i.interest) = toLower('fine dining') " + "AND toLower(p.gender) = toLower('male') " + "AND c.city = 'London' AND c.country = 'United Kingdom' " + "RETURN count(p) AS numPersons" + ) + + result = _mk_graph(nodes, edges).gfql(query, engine=engine) + assert _to_pandas_df(result._nodes).to_dict(orient="records") == [{"numPersons": 4}] + + +@pytest.mark.parametrize("engine", ["pandas", "polars"]) +def test_t1_connected_comma_two_star_fast_path_groups_second_leaf_props(engine: str) -> None: + if engine == "polars": + pytest.importorskip("polars") + nodes = pd.DataFrame({ + "id": [0, 1, 10, 11, 20], + "node_type": ["Person", "Person", "City", "City", "Interest"], + "gender": ["female", "female", None, None, None], + "city": [None, None, "London", "Bristol", None], + "country": [None, None, "United Kingdom", "United Kingdom", None], + "state": [None, None, "England", "England", None], + "interest": [None, None, None, None, "tennis"], + }) + edges = pd.DataFrame({ + "s": [0, 1, 0, 1], + "d": [20, 20, 10, 11], + "rel": ["HAS_INTEREST", "HAS_INTEREST", "LIVES_IN", "LIVES_IN"], + }) + query = ( + "MATCH (p {node_type:'Person'})-[{rel:'HAS_INTEREST'}]->(i {node_type:'Interest'}), " + "(p)-[{rel:'LIVES_IN'}]->(c {node_type:'City'}) " + "WHERE toLower(i.interest) = toLower('tennis') AND toLower(p.gender) = toLower('female') " + "RETURN count(p) AS numPersons, c.city AS city, c.country AS country " + "ORDER BY city" + ) + + result = _mk_graph(nodes, edges).gfql(query, engine=engine) + assert _to_pandas_df(result._nodes).to_dict(orient="records") == [ + {"numPersons": 1, "city": "Bristol", "country": "United Kingdom"}, + {"numPersons": 1, "city": "London", "country": "United Kingdom"}, + ] + + +@pytest.mark.parametrize("engine", ["pandas", "polars"]) +def test_t6_connected_comma_two_star_direct_count_preserves_multiplicity(engine: str) -> None: + if engine == "polars": + pytest.importorskip("polars") + nodes = pd.DataFrame({ + "id": [0, 10, 20], + "node_type": ["Person", "City", "Interest"], + "gender": ["male", None, None], + "city": [None, "London", None], + "country": [None, "United Kingdom", None], + "interest": [None, None, "Fine Dining"], + }) + edges = pd.DataFrame({ + "s": [0, 0, 0, 0], + "d": [20, 20, 10, 10], + "rel": ["HAS_INTEREST", "HAS_INTEREST", "LIVES_IN", "LIVES_IN"], + }) + query = ( + "MATCH (p {node_type:'Person'})-[{rel:'HAS_INTEREST'}]->(i {node_type:'Interest'}), " + "(p)-[{rel:'LIVES_IN'}]->(c {node_type:'City'}) " + "WHERE toLower(i.interest) = toLower('fine dining') " + "AND toLower(p.gender) = toLower('male') " + "AND c.city = 'London' AND c.country = 'United Kingdom' " + "RETURN count(p) AS numPersons" + ) + graph = _mk_graph(nodes, edges) + plan = _compiled_connected_join_plan(query) + + direct = _connected_join_two_star_fast_grouped_count(graph, plan, engine=Engine(engine)) + assert direct is not None + assert _to_pandas_df(direct).to_dict(orient="records") == [{"numPersons": 4}] + + result = graph.gfql(query, engine=engine) + assert _to_pandas_df(result._nodes).to_dict(orient="records") == [{"numPersons": 4}] + + +@pytest.mark.parametrize("engine", ["pandas", "polars"]) +def test_t6_connected_comma_two_star_direct_grouped_count_orders_and_limits(engine: str) -> None: + if engine == "polars": + pytest.importorskip("polars") + nodes = pd.DataFrame({ + "id": [0, 1, 10, 11, 20], + "node_type": ["Person", "Person", "City", "City", "Interest"], + "gender": ["female", "female", None, None, None], + "city": [None, None, "London", "Bristol", None], + "country": [None, None, "United Kingdom", "United Kingdom", None], + "state": [None, None, "England", "England", None], + "interest": [None, None, None, None, "tennis"], + }) + edges = pd.DataFrame({ + "s": [0, 0, 1, 0, 1], + "d": [20, 20, 20, 10, 11], + "rel": ["HAS_INTEREST", "HAS_INTEREST", "HAS_INTEREST", "LIVES_IN", "LIVES_IN"], + }) + query = ( + "MATCH (p {node_type:'Person'})-[{rel:'HAS_INTEREST'}]->(i {node_type:'Interest'}), " + "(p)-[{rel:'LIVES_IN'}]->(c {node_type:'City'}) " + "WHERE toLower(i.interest) = toLower('tennis') AND toLower(p.gender) = toLower('female') " + "RETURN count(p) AS numPersons, c.city AS city, c.country AS country " + "ORDER BY numPersons DESC, city LIMIT 1" + ) + graph = _mk_graph(nodes, edges) + plan = _compiled_connected_join_plan(query) + + direct = _connected_join_two_star_fast_grouped_count(graph, plan, engine=Engine(engine)) + assert direct is not None + assert _to_pandas_df(direct).to_dict(orient="records") == [ + {"city": "London", "country": "United Kingdom", "numPersons": 2} + ] + + result = graph.gfql(query, engine=engine) + assert _to_pandas_df(result._nodes).to_dict(orient="records") == [ + {"city": "London", "country": "United Kingdom", "numPersons": 2} + ] + + +def test_t9_connected_comma_two_star_direct_polars_native_reuses_node_filter_cache() -> None: + pl = pytest.importorskip("polars") + nodes = pd.DataFrame({ + "id": [0, 1, 10, 11, 20], + "node_type": ["Person", "Person", "City", "City", "Interest"], + "gender": ["female", "female", None, None, None], + "city": [None, None, "London", "Bristol", None], + "country": [None, None, "United Kingdom", "United Kingdom", None], + "interest": [None, None, None, None, "tennis"], + }) + edges = pd.DataFrame({ + "s": [0, 0, 1, 0, 1], + "d": [20, 20, 20, 10, 11], + "rel": ["HAS_INTEREST", "HAS_INTEREST", "HAS_INTEREST", "LIVES_IN", "LIVES_IN"], + }) + query = ( + "MATCH (p {node_type:'Person'})-[{rel:'HAS_INTEREST'}]->(i {node_type:'Interest'}), " + "(p)-[{rel:'LIVES_IN'}]->(c {node_type:'City'}) " + "WHERE toLower(i.interest) = toLower('tennis') AND toLower(p.gender) = toLower('female') " + "RETURN count(p) AS numPersons, c.city AS city, c.country AS country " + "ORDER BY numPersons DESC, city LIMIT 1" + ) + graph = cast( + _CypherTestGraph, + _CypherTestGraph().nodes(pl.from_pandas(nodes), "id").edges(pl.from_pandas(edges), "s", "d"), + ) + plan = _compiled_connected_join_plan(query) + + direct = _connected_join_two_star_fast_grouped_count(graph, plan, engine=Engine.POLARS) + assert direct is not None + assert _to_pandas_df(direct).to_dict(orient="records") == [ + {"city": "London", "country": "United Kingdom", "numPersons": 2} + ] + + cache = getattr(graph, "_gfql_connected_join_node_filter_cache") + assert isinstance(cache, dict) + assert len(cache) == 2 + cache_ids = {key: id(value) for key, value in cache.items()} + node_ids_cache = getattr(graph, "_gfql_connected_join_node_ids_cache") + assert isinstance(node_ids_cache, dict) + assert len(node_ids_cache) == 3 + node_ids_cache_ids = {key: id(value) for key, value in node_ids_cache.items()} + first_arm_cache = getattr(graph, "_gfql_connected_join_first_arm_shared_counts_cache") + assert isinstance(first_arm_cache, dict) + assert len(first_arm_cache) == 1 + first_arm_cache_ids = {key: id(value) for key, value in first_arm_cache.items()} + second_arm_cache = getattr(graph, "_gfql_connected_join_second_arm_group_rows_cache") + assert isinstance(second_arm_cache, dict) + assert len(second_arm_cache) == 1 + second_arm_cache_ids = {key: id(value) for key, value in second_arm_cache.items()} + + direct_again = _connected_join_two_star_fast_grouped_count(graph, plan, engine=Engine.POLARS) + assert direct_again is not None + assert _to_pandas_df(direct_again).to_dict(orient="records") == [ + {"city": "London", "country": "United Kingdom", "numPersons": 2} + ] + assert len(cache) == 2 + assert {key: id(value) for key, value in cache.items()} == cache_ids + assert len(node_ids_cache) == 3 + assert {key: id(value) for key, value in node_ids_cache.items()} == node_ids_cache_ids + assert len(first_arm_cache) == 1 + assert {key: id(value) for key, value in first_arm_cache.items()} == first_arm_cache_ids + assert len(second_arm_cache) == 1 + assert {key: id(value) for key, value in second_arm_cache.items()} == second_arm_cache_ids + + +def test_t1_connected_comma_two_star_fast_path_rejects_first_leaf_props() -> None: + query = ( + "MATCH (p {node_type:'Person'})-[{rel:'HAS_INTEREST'}]->(i {node_type:'Interest'}), " + "(p)-[{rel:'LIVES_IN'}]->(c {node_type:'City'}) " + "WHERE toLower(i.interest) = toLower('fine dining') " + "RETURN i.interest AS interest, count(p) AS numPersons" + ) + + plan = _compiled_connected_join_plan(query) + assert _connected_join_two_star_fast_rows(_mk_graph_benchmark_t1_shape_graph(), plan, engine=Engine.PANDAS) is None + + +def test_t1_connected_comma_edge_filter_cache_reuses_simple_edge_match() -> None: + graph = _mk_graph_benchmark_t1_shape_graph() + edges = graph._edges + assert edges is not None + + first = _connected_join_cached_edge_filter(graph, edges, {"rel": "HAS_INTEREST"}, engine=Engine.PANDAS) + second = _connected_join_cached_edge_filter(graph, edges, {"rel": "HAS_INTEREST"}, engine=Engine.PANDAS) + + assert first is second + assert set(first["rel"].unique()) == {"HAS_INTEREST"} + def test_issue_1413_ic3_entity_membership_positive_same_city_friend_only() -> None: graph = _mk_ic3_cross_country_shape_graph() @@ -17643,6 +17918,31 @@ def test_count_table_frame_op_error_and_empty_paths() -> None: out3 = frame_ops.count_table(ctx3, table="nodes", alias="c") assert out3._nodes.to_dict(orient="records") == [{"c": 0}] + +def test_connected_join_post_property_columns_walks_nested_params() -> None: + plan = ConnectedMatchJoinPlan( + pattern_chains=(), + pattern_shared_node_aliases=(), + post_join_chain=Chain( + [ + ASTCall( + "select", + { + "items": [ + ("city", "a.city"), + ("nested", {"left": ["a.country", "b.age", "a.city"]}), + ] + }, + ), + ], + validate=False, + ), + ) + + assert _connected_join_post_property_columns(plan, "a") == ["city", "country"] + assert _connected_join_post_property_columns(plan, "b") == ["age"] + + def test_string_cypher_parenthesized_and_preserves_missing_property_as_null() -> None: # A parenthesized AND must stay on the row-filter path, not the filter_dict path: # the row engine treats an absent property as null (Cypher semantics), whereas