diff --git a/CHANGELOG.md b/CHANGELOG.md index c57afde8b6..7f91071bb5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,7 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm - **GFQL engine-selection docs (pandas / polars / cuDF / polars-gpu)**: New :doc:`Choosing a GFQL Engine ` page — a numbers-first, persona-tested guide to the four interchangeable engines. Adds the one-keyword `engine='polars'` speedup (up to ~38× over pandas on real graphs, no GPU), a motivating warm-median comparison table on real public graphs (LiveJournal 35M / Orkut 117M), a decision matrix (workload shape × size × hardware → engine, with the measured ~10K-edge CPU crossover, GPU-work-bound rule, polars-gpu memory-pressure caveat, and GPU-or-error contract), a cuDF-vs-polars-gpu disambiguation (eager-op vs fused-lazy; cuDF is not deprecated), an honest "when *not* to use Polars" section, the differential-parity guarantee, and a methodology + reproducer-script disclosure. Rewrote the top of `gfql/performance.rst` to lead with the engine comparison (de-marketed the prose), wired the new page into the GFQL toctree + recommended paths, and added Polars/polars-gpu to the engine examples in `gfql/quick.rst` and `gfql/about.rst` (previously only pandas/cuDF were documented). Driven by 4-persona doc user-testing (pandas DS, RAPIDS user, perf engineer, skeptical evaluator). ### Fixed +- **GFQL connected OPTIONAL MATCH no longer crashes on polars frames (engine-polymorphic seed extraction)**: `_optional_arm_start_nodes` (gfql_unified.py) applied pandas-only ops (`.dropna()`/`.drop_duplicates()`/`.rename(columns=)`/boolean-mask indexing) to the joined binding rows, so an IS7-shaped `MATCH ... OPTIONAL MATCH ...` on `engine='polars'` crashed with `AttributeError: 'DataFrame' object has no attribute 'dropna'` before reaching the row pipeline (hit by LDBC SNB interactive-short-7 via the official query text). Now FIXED end-to-end: the seed extraction branches on frame type (`drop_nulls().unique().rename({...})` + `filter(is_in)` on polars), the optional-arm dispatch skips the pandas-only seeded-rows contract on polars (the seed restriction is a pruning hint — arm rows are left-outer-joined on the shared aliases, so unseeded is semantically identical, and `where`-arms already ran unseeded), and the arm join/alias-synthesis block gains a polars twin (`is_in` semi-join prune, `join(how='left')`, null-preserving alias synthesis). Result: **connected OPTIONAL MATCH now runs natively on polars** (previously pandas-only), oracle-exact on the IS7 shape; the simple-CASE null-literal equality (`CASE x WHEN null`, lowered as the `__cypher_case_eq__(x, null)` marker) now lowers natively to `is_null()` on polars (pandas-parity null-mask semantics; the general `CASE x WHEN v` form still declines honestly — it carries pandas' bool/numeric cross-dtype rules), so the full LDBC IS7 query INCLUDING its `CASE r WHEN null` flag projection runs natively on polars, oracle-exact. The polars arm is also PRUNED like the pandas one: since the polars bindings-row path declines `start_nodes` by contract, the same restriction is expressed as an id-membership filter injected into the arm's first (shared) node op — measured 9.2× on the LDBC SF0.1 harness (IS7 message-replies cypher 710.8→77.2 ms). pandas/cuDF paths byte-identical. Regression tests `test_optional_match_polars_frames.py` (pandas oracle, polars native end-to-end, polars no-pandas-ism contract). - **GFQL polars/polars-gpu seeded traversal no longer pays an O(E) NaN re-scan per hop on float-column graphs (identity-stable `_pl_nan_to_null` + clean-cache)**: every `hop()` runs `_coerce_input_formats` → `_pl_nan_to_null`, which scanned `is_nan().any()` over every float column on the (unchanged) resident edge frame **on every call** — so repeated seeded hops on a resident graph (the seeded-Search / native-hop pattern) grew O(E) with edge count on polars while pandas stayed flat. Measured: an indexed polars `g.hop` on 8M edges with one float column was **33.8 ms and growing**; it is now **0.20 ms and FLAT** (~140× faster; O(degree)), matching pandas. `_pl_nan_to_null` now probes an eager polars frame for real NaN once, returns a clean frame UNCHANGED (restoring the #1726 identity guard that #1731 reverted — so frame-identity caches like the #1658 index keep engaging), rewrites only the columns that genuinely carry NaN (values identical to the old unconditional `fill_nan`), and caches the id of frames verified clean so later calls skip the probe (recycle-safe via `weakref.finalize`; a rebound/new frame re-probes, so NaN→null semantics are unchanged). Regression: full 4-engine `test_index.py` parity + new `TestPlNanCleanCache` cases (NaN-present cleaned, clean-frame cached same-object, distinct frames independent). - **GFQL seeded `gfql()`/Cypher chains now engage the resident #1658 adjacency index on ALL FOUR engines (incl. typed edges) instead of always scanning**: a seeded hop expressed as a `gfql([...])`/Cypher CHAIN (rather than a direct `g.hop(...)`) previously fell back to the O(E) scan even with a resident index, because both chain executors attach a synthetic per-edge id column to the edge frame (pandas/cuDF eager: `copy(deep=False)` + assign; polars native lazy: `with_row_index`) → a fresh edge-frame object → the index's `source_ref is df` identity guard missed. Both chain paths now re-point the resident edge adjacency indexes at the augmented frame via `GfqlIndexRegistry.rebind_edges` (provably safe: the shallow copy / `with_row_index` preserves the indexed src/dst columns by value; the structural fingerprint — row count + bound columns + engine — is unchanged by an added column). Additionally, a **simple scalar-equality `edge_match`** (typed edges, e.g. `-[:KNOWS]->` / `e_forward({"type": X})`) is now index-coverable on the wavefront path: `index_seeded_hop` applies the edge predicate to the CSR-matched rows each hop, parity-exact with the scan's `filter_edges_by_dict`. Predicate/membership-list `edge_match`, `edge_query`, and the direct-hop (non-wavefront) path deliberately stay on the scan path (no over-reach). Measured (dgx, 500k nodes / 4M edges / 4 edge types, warm median, index result == scan result): pandas typed 1-hop **2.1×** / untyped **2.7×**; cuDF typed **1.9×** / untyped **2.0×**; polars typed **1.3×** / untyped **5.4×**; polars-gpu typed engaged (1.0×) / untyped **12.4×**; 2-hop typed engaged on all four engines. 29 new differential + engagement regression tests (`test_index.py`), 4-engine, plus the existing 109-test index suite green. - **GFQL Cypher `CASE` with mixed-dtype branches is now engine-consistent on cuDF (no more `GFQLTypeError`)**: pandas coerces the two `CASE` branches to a common type, but cuDF's `.where` raised `TypeError: cudf does not support mixed types`, surfaced as a hard `GFQLTypeError`. This bit `CASE WHEN path IS NULL THEN -1 ELSE length(path) END` over an UNREACHABLE `shortestPath` (int `-1` branch vs an object/null hops branch) — pandas returned `-1`, cuDF errored. The row-AST `CASE` evaluator now unifies the branch dtypes and retries, so cuDF returns the same value pandas does. Regression test `test_case_mixed_dtype.py`. diff --git a/bin/test-polars.sh b/bin/test-polars.sh index d723dca1d3..08cab09be0 100755 --- a/bin/test-polars.sh +++ b/bin/test-polars.sh @@ -24,6 +24,7 @@ POLARS_TEST_FILES=( graphistry/tests/compute/gfql/cypher/test_order_by_null_placement.py graphistry/tests/compute/gfql/test_conformance_ledger.py graphistry/tests/compute/gfql/test_polars_nan_clean.py + graphistry/tests/compute/gfql/test_optional_match_polars_frames.py # index tests exercise the seeded-index hook in the polars hop entry (hop.py) — without # them the hook dominates the now-thin file and trips its per-file coverage floor graphistry/tests/compute/gfql/index/test_index.py diff --git a/graphistry/compute/gfql/lazy/engine/polars/row_pipeline.py b/graphistry/compute/gfql/lazy/engine/polars/row_pipeline.py index c5454a2802..e6917ab3c7 100644 --- a/graphistry/compute/gfql/lazy/engine/polars/row_pipeline.py +++ b/graphistry/compute/gfql/lazy/engine/polars/row_pipeline.py @@ -93,6 +93,21 @@ def _lower_function(node: FunctionCall, columns: Sequence[str]) -> Optional[pl.E parity-verified mappings admitted; anything else returns None (caller NIEs, never guesses).""" import polars as pl # function-local: polars is an optional dependency name = node.name.lower() + if name == "__cypher_case_eq__" and len(node.args) == 2: + # Simple-CASE equality marker (`CASE x WHEN v`). Cypher/pandas semantics: + # null matches null (NOT 3-valued suppressed). Lower ONLY the null-literal + # forms (`CASE x WHEN null` = the LDBC IS7 shape) as the other side's + # null-mask, exactly like the pandas evaluator; the general form carries + # pandas' bool/numeric cross-dtype rules — decline it rather than diverge. + from graphistry.compute.gfql.expr_parser import Literal as _Lit + a_node, b_node = node.args + if isinstance(b_node, _Lit) and b_node.value is None: + a = lower_expr(a_node, columns) + return None if a is None else a.is_null() + if isinstance(a_node, _Lit) and a_node.value is None: + b = lower_expr(b_node, columns) + return None if b is None else b.is_null() + return None args: List[pl.Expr] = [] for arg in node.args: lowered = lower_expr(arg, columns) diff --git a/graphistry/compute/gfql_unified.py b/graphistry/compute/gfql_unified.py index d043698044..d0a6f0603e 100644 --- a/graphistry/compute/gfql_unified.py +++ b/graphistry/compute/gfql_unified.py @@ -6,7 +6,7 @@ from types import MappingProxyType from typing import Any, Dict, List, Literal, Mapping, Optional, Sequence, Set, Tuple, Union, cast from graphistry.Plottable import Plottable -from graphistry.Engine import Engine, EngineAbstract, POLARS_ENGINES, df_concat, df_cons, df_to_engine, df_unique, resolve_engine +from graphistry.Engine import Engine, EngineAbstract, POLARS_ENGINES, df_concat, df_cons, df_to_engine, df_unique, is_polars_df, resolve_engine from graphistry.util import setup_logger from .ast import ASTObject, ASTLet, ASTNode, ASTEdge, ASTCall from .chain import Chain, chain as chain_impl @@ -343,6 +343,51 @@ def _split_binding_and_post_ops(ops: Sequence[ASTObject]) -> Tuple[List[ASTObjec df_ctor = df_cons(concrete_engine) node_col = str(getattr(base_graph, "_node", "id")) + def _optional_arm_membership_chain( + binding_ops: Sequence[ASTObject], + shared_node_aliases: Sequence[str], + joined_rows: DataFrameT, + ) -> Optional[List[ASTObject]]: + """Polars twin of the start_nodes pruning: rewrite the arm's first node op with an + id-membership filter (the shared alias's bound ids) instead of seeding via + start_nodes, which the polars bindings-row path declines by contract. Returns the + pruned binding chain, or None when the shape doesn't qualify (run unseeded).""" + if not binding_ops: + return None + first_op = binding_ops[0] + if not isinstance(first_op, ASTNode): + return None + first_alias = getattr(first_op, "_name", None) + if not isinstance(first_alias, str) or first_alias not in shared_node_aliases: + return None + if first_op.query is not None: + return None + filter_dict = dict(first_op.filter_dict or {}) + if node_col in filter_dict: + return None # id already constrained; don't intersect two conditions on one key + joined_col = next( + ( + col + for col in (f"{first_alias}.{node_col}", first_alias) + if col in joined_rows.columns + ), + None, + ) + if joined_col is None: + return None + seed_col = joined_rows[joined_col] + if is_polars_df(joined_rows): + seed_ids = seed_col.drop_nulls().unique().to_list() + else: + seed_ids = seed_col.dropna().drop_duplicates().tolist() + if not seed_ids: + return None + pruned_first = ASTNode( + filter_dict={**filter_dict, node_col: seed_ids}, + name=first_alias, + ) + return [pruned_first, *binding_ops[1:]] + def _optional_arm_start_nodes( binding_ops: Sequence[ASTObject], shared_node_aliases: Sequence[str], @@ -374,15 +419,16 @@ def _optional_arm_start_nodes( if joined_col is None: return None - seed_frame = cast( - DataFrameT, - df_to_engine( - joined_rows[[joined_col]].dropna().drop_duplicates().rename(columns={joined_col: node_col}), - concrete_engine, - ), - ) + seed_src = joined_rows[[joined_col]] + if is_polars_df(seed_src): + seed_src = seed_src.drop_nulls().unique().rename({joined_col: node_col}) + else: + seed_src = seed_src.dropna().drop_duplicates().rename(columns={joined_col: node_col}) + seed_frame = cast(DataFrameT, df_to_engine(seed_src, concrete_engine)) seed_ids = cast(SeriesT, seed_frame[node_col]) node_ids = cast(SeriesT, base_nodes[node_col]) + if is_polars_df(base_nodes): + return cast(DataFrameT, base_nodes.filter(node_ids.is_in(seed_ids))) return cast(DataFrameT, base_nodes[node_ids.isin(seed_ids)].copy()) # Run base chain to get binding rows. @@ -404,16 +450,33 @@ def _optional_arm_start_nodes( # Chained left-outer-join: one pass per OPTIONAL MATCH arm. for arm in plan.arms: opt_binding_chain, opt_post_ops = _split_binding_and_post_ops(arm.chain.chain) + # The seed restriction is a pruning hint, not semantics: the arm's rows are + # left-outer-joined on the shared aliases, so unpruned arm rows that can't + # join are dropped (`where`-arms already run unseeded). The polars native + # bindings-row path declines seeded runs (start_nodes = bounded-reentry + # contract, pandas-only), so on polars the SAME pruning is expressed as an + # id-membership filter on the arm's first (shared) node op instead. + opt_start_nodes = None + if not arm.chain.where: + if concrete_engine in POLARS_ENGINES: + pruned_chain = _optional_arm_membership_chain( + opt_binding_chain, + arm.shared_node_aliases, + joined, + ) + if pruned_chain is not None: + opt_binding_chain = pruned_chain + else: + opt_start_nodes = _optional_arm_start_nodes( + opt_binding_chain, + arm.shared_node_aliases, + joined, + ) opt_binding_ops = serialize_binding_ops(opt_binding_chain) opt_with_rows = Chain( list(opt_binding_chain) + [ASTCall("rows", {"binding_ops": opt_binding_ops})] + opt_post_ops, where=arm.chain.where, ) - opt_start_nodes = None if arm.chain.where else _optional_arm_start_nodes( - opt_binding_chain, - arm.shared_node_aliases, - joined, - ) opt_rows_result = _chain_dispatch( base_graph, opt_with_rows, @@ -440,7 +503,32 @@ def _optional_arm_start_nodes( and alias in opt_rows_df.columns ] - if opt_rows_df is not None and len(opt_rows_df) > 0 and join_cols: + if is_polars_df(joined): + # polars twin of the pandas block below: same semi-join prune + + # left-outer join + alias synthesis, with null-preserving semantics + # (polars nulls need no NaN->None normalization). + import polars as pl + if opt_rows_df is not None and len(opt_rows_df) > 0 and join_cols: + opt_only_cols = [c for c in opt_rows_df.columns if c not in joined.columns or c in join_cols] + if len(join_cols) == 1: + jc = join_cols[0] + opt_rows_df = opt_rows_df.filter(pl.col(jc).is_in(joined[jc])) + else: + join_keys = joined.select(join_cols).unique() + opt_rows_df = opt_rows_df.join(join_keys, on=join_cols, how="inner") + joined = joined.join(opt_rows_df.select(opt_only_cols), on=join_cols, how="left") + else: + for alias in arm.opt_only_aliases: + if alias not in joined.columns: + joined = joined.with_columns(pl.lit(None).alias(alias)) + for alias in arm.opt_only_aliases: + if alias in joined.columns: + continue + prefix = f"{alias}." + marker_col = next((c for c in joined.columns if c.startswith(prefix)), None) + if marker_col is not None: + joined = joined.with_columns(pl.col(marker_col).alias(alias)) + elif opt_rows_df is not None and len(opt_rows_df) > 0 and join_cols: opt_only_cols = [c for c in opt_rows_df.columns if c not in joined.columns or c in join_cols] # Semi-join filter: restrict opt rows to join-key values present in base # result before materialization. Prevents cross-product blowup when the @@ -457,15 +545,17 @@ def _optional_arm_start_nodes( if alias not in joined.columns: joined[alias] = None - # Synthesize bare alias columns for edge aliases in this arm. - for alias in arm.opt_only_aliases: - if alias in joined.columns: - continue - prefix = f"{alias}." - marker_col = next((c for c in joined.columns if c.startswith(prefix)), None) - if marker_col is not None: - marker = joined[marker_col] - joined[alias] = marker.where(marker.notna(), other=None) + # Synthesize bare alias columns for edge aliases in this arm (pandas/cuDF; + # the polars branch above does its own null-preserving synthesis). + if not is_polars_df(joined): + for alias in arm.opt_only_aliases: + if alias in joined.columns: + continue + prefix = f"{alias}." + marker_col = next((c for c in joined.columns if c.startswith(prefix)), None) + if marker_col is not None: + marker = joined[marker_col] + joined[alias] = marker.where(marker.notna(), other=None) # Delegate RETURN / ORDER BY / SKIP / LIMIT to the standard row pipeline. joined_plottable = base_graph.bind() diff --git a/graphistry/tests/compute/gfql/test_optional_match_polars_frames.py b/graphistry/tests/compute/gfql/test_optional_match_polars_frames.py new file mode 100644 index 0000000000..52e2b256d4 --- /dev/null +++ b/graphistry/tests/compute/gfql/test_optional_match_polars_frames.py @@ -0,0 +1,438 @@ +"""Regression: connected OPTIONAL MATCH seed extraction must be engine-polymorphic. + +_optional_arm_start_nodes (gfql_unified.py) applied pandas-only frame ops +(.dropna()/.drop_duplicates()/.rename(columns=)/boolean-mask __getitem__) to the joined +binding rows, so an IS7-shaped Cypher query (MATCH ... OPTIONAL MATCH ...) on +engine='polars' crashed with AttributeError before reaching the row pipeline +(LDBC SNB interactive-short-7 via the pyg-bench harness). The polars row pipeline may +still honestly decline the query (NotImplementedError, parity-or-error by design), but +it must never crash with a pandas-ism. +""" +from typing import Optional + +import pandas as pd +import pytest + +import graphistry + +try: + import polars as pl + HAS_POLARS = True +except ImportError: + HAS_POLARS = False + + +IS7_SHAPED = """ +MATCH (m:Message {id: $messageId })<-[:REPLY_OF]-(c:Comment)-[:HAS_CREATOR]->(p:Person) + OPTIONAL MATCH (m)-[:HAS_CREATOR]->(a:Person)-[r:KNOWS]-(p) + RETURN c.id AS commentId, + c.creationDate AS commentCreationDate, + p.id AS replyAuthorId, + CASE r + WHEN null THEN false + ELSE true + END AS replyAuthorKnowsOriginalMessageAuthor + ORDER BY commentCreationDate DESC, replyAuthorId +""" + + +def _nodes_pd() -> pd.DataFrame: + return pd.DataFrame({ + "id": [1, 2, 3, 10, 11, 12], + "label__Message": [True, False, False, False, False, False], + "label__Comment": [False, True, True, False, False, False], + "label__Person": [False, False, False, True, True, True], + "creationDate": [100, 200, 300, None, None, None], + }) + + +def _edges_pd() -> pd.DataFrame: + # message creator (12) KNOWS alice (10) but not bob (11) -> discriminating flags + return pd.DataFrame({ + "src": [2, 3, 2, 3, 1, 10, 11, 12], + "dst": [1, 1, 10, 11, 12, 11, 10, 10], + "type": ["REPLY_OF", "REPLY_OF", "HAS_CREATOR", "HAS_CREATOR", "HAS_CREATOR", + "KNOWS", "KNOWS", "KNOWS"], + }) + + +def test_optional_match_pandas_oracle() -> None: + g = graphistry.nodes(_nodes_pd(), "id").edges(_edges_pd(), "src", "dst") + res = g.gfql(IS7_SHAPED, params={"messageId": 1}, engine="pandas") + rows = res._nodes.reset_index(drop=True) + assert list(rows["commentId"]) == [3, 2] + # creator (12) KNOWS alice (10) but not bob (11): discriminating flags + assert list(rows["replyAuthorKnowsOriginalMessageAuthor"]) == [False, True] + + +IS7_SHAPED_NO_CASE = """ +MATCH (m:Message {id: $messageId })<-[:REPLY_OF]-(c:Comment)-[:HAS_CREATOR]->(p:Person) + OPTIONAL MATCH (m)-[:HAS_CREATOR]->(a:Person)-[r:KNOWS]-(p) + RETURN c.id AS commentId, + c.creationDate AS commentCreationDate, + p.id AS replyAuthorId + ORDER BY commentCreationDate DESC, replyAuthorId +""" + + +@pytest.mark.skipif(not HAS_POLARS, reason="polars not installed") +def test_optional_match_polars_native_end_to_end() -> None: + """Connected OPTIONAL MATCH + simple RETURN runs natively on polars, oracle-exact.""" + nodes = pl.from_pandas(_nodes_pd()) + edges = pl.from_pandas(_edges_pd()) + g = graphistry.nodes(nodes, "id").edges(edges, "src", "dst") + res = g.gfql(IS7_SHAPED_NO_CASE, params={"messageId": 1}, engine="polars") + rows = res._nodes.to_pandas() if hasattr(res._nodes, "to_pandas") else res._nodes + rows = rows.reset_index(drop=True) + assert list(rows["commentId"]) == [3, 2] + assert list(rows["replyAuthorId"]) == [11, 10] + + +@pytest.mark.skipif(not HAS_POLARS, reason="polars not installed") +def test_optional_match_polars_no_pandasism_crash() -> None: + """Full IS7 (CASE r WHEN null projection) runs natively on polars, oracle-exact. + + The simple-CASE null-literal equality (`__cypher_case_eq__(x, null)`) lowers to + `is_null()` on polars (pandas-parity null-mask semantics). If a future edit + re-declines it, the honest NIE branch keeps this from crashing dishonestly.""" + nodes = pl.from_pandas(_nodes_pd()) + edges = pl.from_pandas(_edges_pd()) + g = graphistry.nodes(nodes, "id").edges(edges, "src", "dst") + try: + res = g.gfql(IS7_SHAPED, params={"messageId": 1}, engine="polars") + except NotImplementedError: + return # honest parity-or-error decline is acceptable; AttributeError is not + rows = res._nodes.to_pandas() if hasattr(res._nodes, "to_pandas") else res._nodes + rows = rows.reset_index(drop=True) + assert list(rows["commentId"]) == [3, 2] + assert list(rows["replyAuthorKnowsOriginalMessageAuthor"]) == [False, True] + + +# --------------------------------------------------------------------------- +# Amplification round 001 (plans/gfql-optional-match-polars-amplification): +# oracle-gated parity matrix for the polars-native connected OPTIONAL MATCH +# lowering (_optional_arm_membership_chain pruning, polars arm-join twin, +# engine-polymorphic seed extraction). +# +# Two assertion modes: +# - _assert_parity: shape runs natively on polars today; polars must return +# exactly the pandas-oracle rows (regression to NIE or wrong rows fails). +# - _assert_parity_or_nie: shape currently declines honestly on polars +# (parity-or-error contract); a decline must be NotImplementedError, and if +# polars ever starts running the shape it must match the oracle. Silent +# wrong results or a pandas-ism crash always fail. +# +# Fixtures are deliberately asymmetric (distinct score/weight per node/edge) +# so mis-joins produce wrong VALUES, not accidentally-identical rows. +# --------------------------------------------------------------------------- + + +def _amp_nodes_pd() -> pd.DataFrame: + return pd.DataFrame({ + "id": [1, 2, 10, 11, 20, 21], + "label__Message": [True, True, False, False, False, False], + "label__Person": [False, False, True, True, False, False], + "label__Tag": [False, False, False, False, True, True], + "score": [11, 22, 43, 54, 75, 86], # unique per node: discriminating + }) + + +def _amp_edges_pd() -> pd.DataFrame: + # msg1 -HAS_CREATOR-> alice(10) -LIKES-> tag20 + # msg2 -HAS_CREATOR-> bob(11) (bob likes nothing: unmatched seed) + # msg1 -HAS_TAG-> tag20, msg2 -HAS_TAG-> tag21 + return pd.DataFrame({ + "src": [1, 2, 10, 1, 2], + "dst": [10, 11, 20, 20, 21], + "type": ["HAS_CREATOR", "HAS_CREATOR", "LIKES", "HAS_TAG", "HAS_TAG"], + "weight": [1.5, 2.5, 3.5, 4.5, 5.5], # unique per edge: discriminating + }) + + +def _run_engine(nodes: pd.DataFrame, edges: pd.DataFrame, query: str, engine: str) -> pd.DataFrame: + if engine == "polars": + g = graphistry.nodes(pl.from_pandas(nodes), "id").edges(pl.from_pandas(edges), "src", "dst") + else: + g = graphistry.nodes(nodes, "id").edges(edges, "src", "dst") + res = g.gfql(query, engine=engine) + out = res._nodes + if hasattr(out, "to_pandas"): + out = out.to_pandas() + return out.reset_index(drop=True) + + +def _normalize(df: pd.DataFrame) -> pd.DataFrame: + """Sorted row set with dtype normalization (int/bool/float -> float64, NaN==null).""" + df = df.copy() + for c in df.columns: + if pd.api.types.is_numeric_dtype(df[c]) or pd.api.types.is_bool_dtype(df[c]): + df[c] = df[c].astype("float64") + elif df[c].isna().all(): + # all-null column: unify to float64 NaN regardless of source dtype + df[c] = df[c].astype("float64") + else: + # unify null representation (polars->pandas gives None, pandas gives NaN) + df[c] = df[c].astype(object).where(df[c].notna(), None) + cols = sorted(df.columns) + return df[cols].sort_values(cols, na_position="last").reset_index(drop=True) + + +def _assert_parity( + nodes: pd.DataFrame, + edges: pd.DataFrame, + query: str, + expected: "Optional[pd.DataFrame]" = None, +) -> None: + oracle = _run_engine(nodes, edges, query, "pandas") + if expected is not None: + pd.testing.assert_frame_equal( + _normalize(oracle[expected.columns.tolist()]), _normalize(expected), + check_dtype=False, + ) + got = _run_engine(nodes, edges, query, "polars") + assert sorted(got.columns) == sorted(oracle.columns) + pd.testing.assert_frame_equal(_normalize(got), _normalize(oracle), check_dtype=False) + + +def _assert_parity_or_nie(nodes: pd.DataFrame, edges: pd.DataFrame, query: str) -> None: + oracle = _run_engine(nodes, edges, query, "pandas") + try: + got = _run_engine(nodes, edges, query, "polars") + except NotImplementedError: + return # honest parity-or-error decline; any other exception fails the test + assert sorted(got.columns) == sorted(oracle.columns) + pd.testing.assert_frame_equal(_normalize(got), _normalize(oracle), check_dtype=False) + + +Q_ARM_LIKES = """ +MATCH (m:Message)-[:HAS_CREATOR]->(p:Person) +OPTIONAL MATCH (p)-[r:LIKES]->(t:Tag) +RETURN m.id AS mid, p.id AS pid, t.id AS tid, t.score AS tscore +ORDER BY mid +""" + +pytestmark_polars = pytest.mark.skipif(not HAS_POLARS, reason="polars not installed") + + +@pytestmark_polars +def test_amp_partial_match_mixed_seeds() -> None: + """Mix of matched (alice->tag20) and unmatched (bob) seeds: null-extension + row must carry nulls, matched row the RIGHT tag's value (score 75).""" + expected = pd.DataFrame({ + "mid": [1, 2], "pid": [10, 11], + "tid": [20.0, None], "tscore": [75.0, None], + }) + _assert_parity(_amp_nodes_pd(), _amp_edges_pd(), Q_ARM_LIKES, expected) + + +@pytestmark_polars +def test_amp_zero_match_arm_case_flag() -> None: + """Arm matches nothing anywhere (no NOPE edges): all seeds get flag False + via the CASE r WHEN null -> is_null lowering.""" + q = """ + MATCH (m:Message)-[:HAS_CREATOR]->(p:Person) + OPTIONAL MATCH (p)-[r:NOPE]->(t:Tag) + RETURN m.id AS mid, p.id AS pid, + CASE r WHEN null THEN false ELSE true END AS hasR + ORDER BY mid + """ + expected = pd.DataFrame({"mid": [1, 2], "pid": [10, 11], "hasR": [False, False]}) + _assert_parity(_amp_nodes_pd(), _amp_edges_pd(), q, expected) + + +@pytestmark_polars +def test_amp_zero_match_arm_property_projection() -> None: + """Zero-match arm + projecting t.id: pandas emits all-null column; polars + currently declines (NIE on 'select': no t.* columns to project).""" + q = """ + MATCH (m:Message)-[:HAS_CREATOR]->(p:Person) + OPTIONAL MATCH (p)-[r:NOPE]->(t:Tag) + RETURN m.id AS mid, p.id AS pid, t.id AS tid + ORDER BY mid + """ + _assert_parity_or_nie(_amp_nodes_pd(), _amp_edges_pd(), q) + + +@pytestmark_polars +def test_amp_multiple_optional_arms() -> None: + """Two OPTIONAL MATCH arms; polars currently declines the second arm + (NIE on 'rows'); must never silently mis-join.""" + q = """ + MATCH (m:Message)-[:HAS_CREATOR]->(p:Person) + OPTIONAL MATCH (p)-[r:LIKES]->(t:Tag) + OPTIONAL MATCH (m)-[h:HAS_TAG]->(t2:Tag) + RETURN m.id AS mid, p.id AS pid, t.id AS tid, t2.id AS t2id + ORDER BY mid + """ + _assert_parity_or_nie(_amp_nodes_pd(), _amp_edges_pd(), q) + + +@pytestmark_polars +def test_amp_arm_where_keeps_rows() -> None: + """WHERE on the arm alias that keeps rows (score>=70 keeps tag20). WHERE + arms skip seeding/pruning entirely (run unseeded) — exercises that branch.""" + q = """ + MATCH (m:Message)-[:HAS_CREATOR]->(p:Person) + OPTIONAL MATCH (p)-[r:LIKES]->(t:Tag) + WHERE t.score >= 70 + RETURN m.id AS mid, p.id AS pid, t.id AS tid, t.score AS tscore + ORDER BY mid + """ + expected = pd.DataFrame({ + "mid": [1, 2], "pid": [10, 11], + "tid": [20.0, None], "tscore": [75.0, None], + }) + _assert_parity(_amp_nodes_pd(), _amp_edges_pd(), q, expected) + + +@pytestmark_polars +def test_amp_arm_where_filters_all_rows() -> None: + """WHERE filtering the arm to zero rows: pandas null-extends every seed; + polars currently declines (NIE) — must not fabricate matched rows.""" + q = """ + MATCH (m:Message)-[:HAS_CREATOR]->(p:Person) + OPTIONAL MATCH (p)-[r:LIKES]->(t:Tag) + WHERE t.score > 9000 + RETURN m.id AS mid, p.id AS pid, t.id AS tid + ORDER BY mid + """ + _assert_parity_or_nie(_amp_nodes_pd(), _amp_edges_pd(), q) + + +@pytestmark_polars +def test_amp_duplicate_seed_ids() -> None: + """Both messages share one creator (alice): the shared-alias join key is + duplicated across base rows; each row must still get alice's arm match.""" + edges = _amp_edges_pd() + edges.loc[1, "dst"] = 10 # msg2 also HAS_CREATOR -> alice + expected = pd.DataFrame({ + "mid": [1, 2], "pid": [10, 10], + "tid": [20.0, 20.0], "tscore": [75.0, 75.0], + }) + _assert_parity(_amp_nodes_pd(), edges, Q_ARM_LIKES, expected) + + +@pytestmark_polars +def test_amp_empty_seed_set() -> None: + """Main MATCH matches nothing (id 999): empty result, no crash in seed + extraction/pruning over an empty joined frame.""" + q = """ + MATCH (m:Message {id: 999})-[:HAS_CREATOR]->(p:Person) + OPTIONAL MATCH (p)-[r:LIKES]->(t:Tag) + RETURN m.id AS mid, t.id AS tid + """ + oracle = _run_engine(_amp_nodes_pd(), _amp_edges_pd(), q, "pandas") + got = _run_engine(_amp_nodes_pd(), _amp_edges_pd(), q, "polars") + assert len(oracle) == 0 + assert len(got) == 0 + + +@pytestmark_polars +def test_amp_arm_fanout_multiplicity() -> None: + """One seed matches two arm rows (alice likes tag20 AND tag21): left join + must fan out to 3 rows total with per-tag values, not dedupe or cross.""" + edges = pd.concat([_amp_edges_pd(), pd.DataFrame({ + "src": [10], "dst": [21], "type": ["LIKES"], "weight": [9.5], + })], ignore_index=True) + q = """ + MATCH (m:Message)-[:HAS_CREATOR]->(p:Person) + OPTIONAL MATCH (p)-[r:LIKES]->(t:Tag) + RETURN m.id AS mid, p.id AS pid, t.id AS tid, t.score AS tscore + ORDER BY mid, tid + """ + expected = pd.DataFrame({ + "mid": [1, 1, 2], "pid": [10, 10, 11], + "tid": [20.0, 21.0, None], "tscore": [75.0, 86.0, None], + }) + _assert_parity(_amp_nodes_pd(), edges, q, expected) + + +@pytestmark_polars +def test_amp_null_carrying_arm_value_column() -> None: + """Matched arm row whose projected value is itself null (tag20.score=null): + null-from-match must survive the join identically to null-from-no-match.""" + nodes = _amp_nodes_pd() + nodes["score"] = nodes["score"].astype("float64") + nodes.loc[nodes["id"] == 20, "score"] = None + expected = pd.DataFrame({ + "mid": [1, 2], "pid": [10, 11], + "tid": [20.0, None], "tscore": [None, None], + }) + _assert_parity(nodes, _amp_edges_pd(), Q_ARM_LIKES, expected) + + +@pytestmark_polars +def test_amp_string_ids() -> None: + """String node ids: membership-pruning filter and joins on Utf8 keys.""" + nodes = _amp_nodes_pd() + nodes["id"] = "n" + nodes["id"].astype(str) + edges = _amp_edges_pd() + edges["src"] = "n" + edges["src"].astype(str) + edges["dst"] = "n" + edges["dst"].astype(str) + expected = pd.DataFrame({ + "mid": ["n1", "n2"], "pid": ["n10", "n11"], + "tid": ["n20", None], "tscore": [75.0, None], + }) + _assert_parity(nodes, edges, Q_ARM_LIKES, expected) + + +@pytestmark_polars +def test_amp_float_join_key_dtype_divergence() -> None: + """Float edge endpoints vs int node ids, and all-float ids: dtype-divergent + join keys must give parity or an honest decline (currently NIE on polars), + never a silently-empty or mis-typed join.""" + edges_f = _amp_edges_pd() + edges_f["src"] = edges_f["src"].astype("float64") + edges_f["dst"] = edges_f["dst"].astype("float64") + _assert_parity_or_nie(_amp_nodes_pd(), edges_f, Q_ARM_LIKES) + + nodes_f = _amp_nodes_pd() + nodes_f["id"] = nodes_f["id"].astype("float64") + _assert_parity_or_nie(nodes_f, edges_f, Q_ARM_LIKES) + + +@pytestmark_polars +def test_amp_pruning_fallback_id_constrained_first_op() -> None: + """Arm's first node op already carries an id constraint ({id: 10}): + _optional_arm_membership_chain must decline (no double-constraint on the id + key) and the unseeded fallback must still produce oracle-exact rows — + alice's row matched, bob's row null-extended.""" + q = """ + MATCH (m:Message)-[:HAS_CREATOR]->(p:Person) + OPTIONAL MATCH (p {id: 10})-[r:LIKES]->(t:Tag) + RETURN m.id AS mid, p.id AS pid, t.id AS tid + ORDER BY mid + """ + expected = pd.DataFrame({ + "mid": [1, 2], "pid": [10, 11], "tid": [20.0, None], + }) + _assert_parity(_amp_nodes_pd(), _amp_edges_pd(), q, expected) + + +@pytestmark_polars +def test_amp_pruning_large_arm_space() -> None: + """Arm alias space >> seed space (120 extra persons each liking a decoy + tag with a distinct score; only 2 seeds): the first-alias id-membership + pruning is logically engaged and must not change results — decoy persons' + arm rows must never leak into the join.""" + n_extra = 120 + extra_ids = list(range(1000, 1000 + n_extra)) + decoy_tag = 5000 + nodes = pd.concat([_amp_nodes_pd(), pd.DataFrame({ + "id": extra_ids + [decoy_tag], + "label__Message": False, + "label__Person": [True] * n_extra + [False], + "label__Tag": [False] * n_extra + [True], + "score": list(range(2000, 2000 + n_extra)) + [9999], + })], ignore_index=True) + edges = pd.concat([_amp_edges_pd(), pd.DataFrame({ + "src": extra_ids, + "dst": decoy_tag, + "type": "LIKES", + "weight": [float(i) for i in range(n_extra)], + })], ignore_index=True) + expected = pd.DataFrame({ + "mid": [1, 2], "pid": [10, 11], + "tid": [20.0, None], "tscore": [75.0, None], + }) + _assert_parity(nodes, edges, Q_ARM_LIKES, expected)