feat(gfql): Polars-GPU engine (engine='polars-gpu', cudf_polars) [PR3, stacked on #1648]

CHANGELOG.md

@@ -11,6 +11,7 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 ### Added
 - **GFQL native Polars engine — traversals (`engine='polars'`)**: Added a native, vectorized Polars execution engine for the core GFQL traversals `hop()` and `chain()`, dispatched at the engine boundary so the production pandas/cuDF paths are untouched. `Engine.POLARS` is opt-in (explicit `engine='polars'`); `engine='auto'` with Polars input still coerces to pandas as before. Covers forward/reverse/undirected single-hop traversal, directed multi-hop chains, node/edge filter dicts and predicates (lowered to Polars expressions), `edge_match`/`source_node_match`/`destination_node_match`, `target_wave_front`, and alias names; the BFS advances via semi/anti joins (no per-row Python work). Validated by differential parity against the pandas engine (hop + chain test suites plus a randomized fuzzer) and benchmarked vs pandas (`benchmarks/gfql/pandas_vs_polars.py`) — Polars wins at scale (up to ~2.5x on multi-edge chains at millions of edges; crossover ~50–100k rows). Variable-length/multi-hop edges, undirected edges in multi-edge chains, hop labels, and node `query=` raise `NotImplementedError` for now (use `engine='pandas'`).
 - **GFQL native Polars engine — cypher row pipeline (`engine='polars'`)**: Extended the Polars engine to the Cypher `MATCH … RETURN` row surface, natively vectorized. **NO CHEATING:** the polars engine never silently falls back to the pandas engine — every query runs natively on polars or raises an honest `NotImplementedError` pointing at `engine='pandas'` (falling back to pandas would misrepresent pandas performance as polars; only a human may consent to a bridge). `chain_polars` splits boundary `call()` ops (mirroring the pandas `_handle_boundary_calls`) and runs each trailing row op per-op native or raises. **Native polars** (no pandas round-trip): frame ops (`rows`/`limit`/`skip`/`distinct`/`drop_cols`), `select`/`with_`/`return_` projection (a conservative cypher-expr-AST → `pl.Expr` lowering covering property access, arithmetic, comparison, boolean, literals), `order_by` (`.sort`), `group_by` (`count`/`sum`/`avg`/`min`/`max`), `unwind` (literal-list cross-join), the result projection for property/expr columns, and entity-text `RETURN n` rendering for int/string/bool nodes (`pl.concat_str`). **Honestly deferred** (raise `NotImplementedError`, no pandas fallback): multi-entity `rows(binding_ops=…)`, cross-entity same-path `WHERE` (`DFSamePathExecutor`), float/temporal/nested entity-text, and exotic expressions (CASE/list/map/temporal, `collect` aggregates) — these are the forward native-engineering targets. Validated by differential parity vs pandas including a TCK-style conformance lane (`test_engine_polars_cypher_conformance.py`: native-only curated corpus + seeded fuzzer + NULL/3-valued-logic graph + entity-text escaping, plus a `DEFERRED` list asserting deferred queries raise rather than silently bridge) and benchmarked (`benchmarks/gfql/cypher_row_pipeline.py`). **Perf (interleaved, 1M nodes, each engine on its native-frame graph, all fully native):** polars wins **5.6–38×** across the surface — `RETURN n` ~38×, `ORDER BY` ~17×, `WHERE`+`ORDER BY`+`LIMIT` ~14×, traversals 6–7.5×, projections/aggregations/`DISTINCT` 5.6–6.9×. cuDF/pandas paths untouched.
+- **GFQL Polars-GPU engine (`engine='polars-gpu'`, cudf_polars)**: Added a GPU execution mode of the native Polars engine — the same vectorized ops, but the hot traversal joins materialize on GPU via the RAPIDS cudf_polars backend (`LazyFrame.collect(engine=pl.GPUEngine(raise_on_fail=False))`). `Engine.POLARS_GPU` is **explicit opt-in only** (`engine='polars-gpu'`); `engine='auto'` never selects it. Frames stay `pl.DataFrame` (handled exactly like `POLARS` in all frame ops) — only the collect boundary changes; GPU intent is carried by a context var set at the dispatch boundary, so `engine='polars'` (CPU) behavior is byte-for-byte unchanged (when GPU is inactive the join helper is the ordinary eager join). `raise_on_fail=False` keeps any GPU-incapable node on CPU **in Polars** — NOT a pandas bridge (still honest, still native; see NO-CHEATING). Validated by differential parity `engine='polars-gpu' == engine='polars'` across the cypher conformance corpus + core traversals (`test_engine_polars_gpu.py`, skips when no cudf_polars/GPU). GPU-routes the hop semi-joins + final edge/node materialization AND the cypher row-pipeline ops (select / where_rows / order_by / group_by.agg / unwind cross-join) — each helper runs the op lazily and collects on GPU when active, or the ordinary eager op when not. Entity-text rendering remains on CPU-Polars for now (still correct, still native). Benchmarks (cudf_polars spikes) show polars-GPU is fastest on traversal-heavy workloads at 1M–50M edges.
 
 ### Changed
 - **GFQL Cypher parse memoization (perf)**: `parse_cypher` now memoizes its result (LRU over the deterministic lark parse+transform → immutable frozen AST). Repeated identical Cypher queries skip the ~15 ms parse — the dominant per-call cost of small queries (~50% of a Cypher call at 100k rows) — making end-to-end query latency ~1.3–1.7× faster at small/interactive sizes across pandas/polars/cuDF. Safe to share the cached AST: every Cypher AST node is `@dataclass(frozen=True)` and `compile_cypher_query` does not mutate the parsed tree; validation errors still raise and are not cached.

graphistry/Engine.py

@@ -14,19 +14,29 @@ class Engine(Enum):
     DASK = 'dask'
     DASK_CUDF = 'dask_cudf'
     POLARS = 'polars'
+    # GPU execution MODE of the Polars engine (cudf_polars): frames are still
+    # ``pl.DataFrame`` (handled exactly like POLARS in all frame ops); only the
+    # engine's hot materializations collect on GPU. Explicit opt-in only — AUTO
+    # never selects it.
+    POLARS_GPU = 'polars-gpu'
 
 class EngineAbstract(Enum):
     PANDAS = Engine.PANDAS.value
     CUDF = Engine.CUDF.value
     DASK = Engine.DASK.value
     DASK_CUDF = Engine.DASK_CUDF.value
     POLARS = Engine.POLARS.value
+    POLARS_GPU = Engine.POLARS_GPU.value
     AUTO = 'auto'
 
 
 # Type alias for engine parameter - accepts both enum values and string literals
 # Includes 'auto' for automatic detection
-EngineAbstractType = Union[EngineAbstract, Literal['pandas', 'cudf', 'dask', 'dask_cudf', 'polars', 'auto']]
+EngineAbstractType = Union[EngineAbstract, Literal['pandas', 'cudf', 'dask', 'dask_cudf', 'polars', 'polars-gpu', 'auto']]
+
+
+# Engines whose native frame type is Polars (POLARS + its GPU execution mode).
+POLARS_ENGINES = (Engine.POLARS, Engine.POLARS_GPU)
 
 DataframeLike = Any  # pdf, cudf, ddf, dgdf
 DataframeLocalLike = Any  # pdf, cudf
@@ -211,7 +221,7 @@ def df_to_engine(df, engine: Engine):
         if not isinstance(df, pd.DataFrame):
             df = df_to_engine(df, Engine.PANDAS)
         return dd.from_pandas(df, npartitions=1)
-    elif engine == Engine.POLARS:
+    elif engine in POLARS_ENGINES:
         import polars as pl
         if isinstance(df, pl.DataFrame):
             return df
@@ -251,7 +261,7 @@ def df_concat(engine: Engine):
     elif engine == Engine.CUDF:
         import cudf
         return cudf.concat
-    elif engine == Engine.POLARS:
+    elif engine in POLARS_ENGINES:
         return _pl_concat
     elif engine == Engine.DASK:
         raise NotImplementedError("DASK is an input format, not a compute engine — use engine='auto' or engine='pandas'")
@@ -326,7 +336,7 @@ def df_cons(engine: Engine):
     elif engine == Engine.CUDF:
         import cudf
         return cudf.DataFrame
-    elif engine == Engine.POLARS:
+    elif engine in POLARS_ENGINES:
         import polars as pl
         return pl.DataFrame
     raise ValueError(f'Only engines pandas/cudf supported, got: {engine}')
@@ -337,7 +347,7 @@ def s_cons(engine: Engine):
     elif engine == Engine.CUDF:
         import cudf
         return cudf.Series
-    elif engine == Engine.POLARS:
+    elif engine in POLARS_ENGINES:
         import polars as pl
         return pl.Series
     raise ValueError(f'Only engines pandas/cudf supported, got: {engine}')

graphistry/compute/ComputeMixin.py

@@ -86,7 +86,7 @@ def _is_already_correct(df: Any) -> bool:
             if not has_cudf:
                 return True  # cudf unavailable — skip coercion; downstream handles gracefully
             return 'cudf' in type_mod
-        elif engine == Engine.POLARS:
+        elif engine in (Engine.POLARS, Engine.POLARS_GPU):
             return 'polars' in type_mod
         return True
 

graphistry/compute/chain.py

@@ -711,17 +711,20 @@ def chain(
     engine_concrete_early = resolve_engine(engine, self)
     self = _coerce_input_formats(self, engine_concrete_early)
 
-    if engine_concrete_early == Engine.POLARS:
+    if engine_concrete_early in (Engine.POLARS, Engine.POLARS_GPU):
         # Native polars chain lives in a dedicated dispatched module so the
         # production pandas/cuDF orchestration below stays untouched (see
         # plans/gfql-polars-engine). Correctness gated by differential parity.
+        # POLARS_GPU runs the SAME native ops but collects hot joins on GPU.
         if validate_schema:
             Chain(ops if not isinstance(ops, Chain) else ops.chain).validate(collect_all=False)
         from graphistry.compute.gfql.engine_polars.chain import chain_polars
+        from graphistry.compute.gfql.engine_polars.gpu import gpu_mode
         # NO pandas fallback here (see plan.md NO-CHEATING): chain_polars raises
         # NotImplementedError for deferred features (var-length/multi-hop edges,
         # undirected multi-edge); that honest signal propagates to the caller.
-        return chain_polars(self, ops, start_nodes=start_nodes)
+        with gpu_mode(engine_concrete_early == Engine.POLARS_GPU):
+            return chain_polars(self, ops, start_nodes=start_nodes)
 
     if policy:
         from graphistry.compute.gfql.call.executor import _thread_local as call_thread_local

graphistry/compute/gfql/engine_polars/gpu.py

@@ -0,0 +1,89 @@
+"""GPU execution mode for the native Polars engine (``Engine.POLARS_GPU``).
+
+The Polars engine is eager (operates on ``pl.DataFrame``). GPU mode runs the
+SAME native ops but materializes the heavy ones on GPU via the cudf_polars
+backend: ``LazyFrame.collect(engine=pl.GPUEngine(raise_on_fail=False))``.
+
+``raise_on_fail=False`` keeps any GPU-incapable node on CPU **in Polars** — NOT a
+pandas bridge (still honest, still native Polars; see plan.md NO-CHEATING).
+
+GPU intent is carried by a context var set at the dispatch boundary (so the
+engine internals don't need a ``gpu`` parameter threaded through every call).
+When GPU is inactive, helpers run the ordinary eager path verbatim, so the
+``engine='polars'`` (CPU) behavior is byte-for-byte unchanged.
+"""
+from __future__ import annotations
+
+import contextvars
+from typing import Any, Optional
+
+_GPU_ACTIVE: "contextvars.ContextVar[bool]" = contextvars.ContextVar("gfql_polars_gpu_active", default=False)
+
+
+def gpu_active() -> bool:
+    return _GPU_ACTIVE.get()
+
+
+class gpu_mode:
+    """Context manager: activate GPU collection for the enclosed Polars engine run."""
+
+    def __init__(self, active: bool = True) -> None:
+        self._active = active
+        self._token: Optional[contextvars.Token] = None
+
+    def __enter__(self) -> "gpu_mode":
+        self._token = _GPU_ACTIVE.set(self._active)
+        return self
+
+    def __exit__(self, *exc: Any) -> None:
+        if self._token is not None:
+            _GPU_ACTIVE.reset(self._token)
+
+
+def _gpu_engine() -> Any:
+    import polars as pl
+    return pl.GPUEngine(raise_on_fail=False)
+
+
+def collect(lf: Any) -> Any:
+    """Collect a polars LazyFrame — on GPU when GPU mode is active, else CPU."""
+    if gpu_active():
+        return lf.collect(engine=_gpu_engine())
+    return lf.collect()
+
+
+def join(left: Any, right: Any, **kwargs: Any) -> Any:
+    """Eager-signature join that runs on GPU (lazy + GPU collect) when active.
+
+    When GPU is inactive this is exactly ``left.join(right, **kwargs)`` (eager),
+    so CPU behavior is unchanged. ``left``/``right`` are eager ``pl.DataFrame``.
+    """
+    if gpu_active():
+        return left.lazy().join(right.lazy(), **kwargs).collect(engine=_gpu_engine())
+    return left.join(right, **kwargs)
+
+
+# Row-pipeline op helpers: each runs the op on GPU (lazy + GPU collect) when GPU
+# mode is active, else the ordinary eager op verbatim (CPU behavior unchanged).
+def select(df: Any, exprs: Any) -> Any:
+    if gpu_active():
+        return df.lazy().select(exprs).collect(engine=_gpu_engine())
+    return df.select(exprs)
+
+
+def where(df: Any, predicate: Any) -> Any:  # filter (named to avoid shadowing builtin)
+    if gpu_active():
+        return df.lazy().filter(predicate).collect(engine=_gpu_engine())
+    return df.filter(predicate)
+
+
+def sort(df: Any, by: Any, **kwargs: Any) -> Any:
+    if gpu_active():
+        return df.lazy().sort(by, **kwargs).collect(engine=_gpu_engine())
+    return df.sort(by, **kwargs)
+
+
+def group_agg(df: Any, keys: Any, aggs: Any, **kwargs: Any) -> Any:
+    if gpu_active():
+        return df.lazy().group_by(keys, **kwargs).agg(aggs).collect(engine=_gpu_engine())
+    return df.group_by(keys, **kwargs).agg(aggs)

graphistry/compute/gfql/engine_polars/hop.py

@@ -72,6 +72,7 @@ def hop_polars(
 ) -> Plottable:
     import polars as pl
     from graphistry.Engine import Engine, df_to_engine
+    from .gpu import join as _gjoin  # GPU-collects when POLARS_GPU active, else eager join
 
     if direction not in ("forward", "reverse", "undirected"):
         raise ValueError(
@@ -184,7 +185,7 @@ def _idframe(df, col) -> "pl.DataFrame":
         current_hop += 1
 
         frontier_iter = frontier if allowed_source is None else frontier.join(allowed_source, on=NID, how="semi")
-        hop_edges = pairs.join(frontier_iter.rename({NID: FROM}), on=FROM, how="semi")
+        hop_edges = _gjoin(pairs, frontier_iter.rename({NID: FROM}), on=FROM, how="semi")
 
         is_last = not to_fixed_point and resolved_max_hops is not None and current_hop >= resolved_max_hops
         if target_final is not None and is_last:
@@ -208,7 +209,7 @@ def _idframe(df, col) -> "pl.DataFrame":
     else:
         visited_edges = edges_idx.select(pl.col(EID)).clear()
 
-    out_edges = edges_idx.join(visited_edges, on=EID, how="semi")
+    out_edges = _gjoin(edges_idx, visited_edges, on=EID, how="semi")
     if synth_eid:
         out_edges = out_edges.drop(EID)
 
@@ -223,6 +224,6 @@ def _idframe(df, col) -> "pl.DataFrame":
         ).unique(subset=[NID])
         needed = pl.concat([needed, endpoints], how="vertical_relaxed").unique(subset=[NID])
 
-    out_nodes = all_nodes.join(needed.rename({NID: node_col}), on=node_col, how="semi")
+    out_nodes = _gjoin(all_nodes, needed.rename({NID: node_col}), on=node_col, how="semi")
 
     return g.nodes(out_nodes, node_col).edges(out_edges, src, dst)

graphistry/compute/gfql/engine_polars/row_pipeline.py

@@ -15,6 +15,7 @@
 from typing import Any, List, Optional, Sequence, Tuple
 
 from graphistry.Plottable import Plottable
+from . import gpu  # GPU-collects these ops when POLARS_GPU active, else eager (CPU unchanged)
 
 
 def _parser():
@@ -215,7 +216,7 @@ def select_polars(g: Plottable, items: Sequence[Any]) -> Optional[Plottable]:
     exprs = lower_select_items(items, list(table.columns))
     if exprs is None:
         return None
-    return _rewrap(g, table.select(exprs))
+    return _rewrap(g, gpu.select(table, exprs))
 
 
 def where_rows_polars(
@@ -252,7 +253,7 @@ def where_rows_polars(
     combined = preds[0]
     for pred in preds[1:]:
         combined = combined & pred
-    return _rewrap(g, table.filter(combined))
+    return _rewrap(g, gpu.where(table, combined))
 
 
 def order_by_polars(g: Plottable, keys: Sequence[Any]) -> Optional[Plottable]:
@@ -265,7 +266,7 @@ def order_by_polars(g: Plottable, keys: Sequence[Any]) -> Optional[Plottable]:
     # nulls_last=False matches pandas sort_values default (NaN last only for asc);
     # cypher ORDER BY puts NULLs last — polars default is nulls_last=False, so set
     # it explicitly to match the pandas engine's na_position='last'.
-    return _rewrap(g, table.sort(exprs, descending=descending, nulls_last=True))
+    return _rewrap(g, gpu.sort(table, exprs, descending=descending, nulls_last=True))
 
 
 # Aggregation funcs lowered to native polars; collect/collect_distinct/stdev/
@@ -313,7 +314,7 @@ def group_by_polars(g: Plottable, keys: Sequence[Any], aggregations: Sequence[An
         if lowered is None:
             return None
         aggs.append(lowered)
-    out = table.group_by(list(keys), maintain_order=True).agg(aggs)
+    out = gpu.group_agg(table, list(keys), aggs, maintain_order=True)
     return _rewrap(g, out)
 
 
@@ -343,7 +344,7 @@ def unwind_polars(g: Plottable, expr: str, as_: str = "value") -> Optional[Plott
         return None
     values = [it.value for it in node.items if isinstance(it, Literal)]
     rhs = pl.DataFrame({as_: values})
-    return _rewrap(g, table.join(rhs, how="cross"))
+    return _rewrap(g, gpu.join(table, rhs, how="cross"))
 
 
 def can_select_native(items: Sequence[Any], columns: Sequence[str]) -> bool:

graphistry/compute/gfql_unified.py

@@ -1657,7 +1657,10 @@ def _chain_dispatch(
     context: ExecutionContext,
     start_nodes: Optional[DataFrameT] = None,
 ) -> Plottable:
-    if chain_obj.where and engine in (EngineAbstract.POLARS, "polars", Engine.POLARS):
+    if chain_obj.where and engine in (
+        EngineAbstract.POLARS, "polars", Engine.POLARS,
+        EngineAbstract.POLARS_GPU, "polars-gpu", Engine.POLARS_GPU,
+    ):
         # Cross-entity / same-path WHERE routes through DFSamePathExecutor
         # (df_executor.py), which has no native polars implementation. NO pandas
         # fallback (see plan.md NO-CHEATING) — raise honestly.

graphistry/compute/hop.py

@@ -109,23 +109,26 @@ def hop(self: Plottable,
     from graphistry.compute.ComputeMixin import _coerce_input_formats  # lazy — avoids circular import
     self = _coerce_input_formats(self, engine_concrete)
 
-    if engine_concrete == Engine.POLARS:
+    if engine_concrete in (Engine.POLARS, Engine.POLARS_GPU):
         # Native polars traversal lives in a dedicated dispatched module so the
         # production pandas/cuDF internals below stay untouched (see
         # plans/gfql-polars-engine). Correctness gated by differential parity.
+        # POLARS_GPU runs the SAME native ops but collects the hot joins on GPU.
         from graphistry.compute.gfql.engine_polars import hop_polars
-        return hop_polars(
-            self, nodes, hops,
-            min_hops=min_hops, max_hops=max_hops,
-            output_min_hops=output_min_hops, output_max_hops=output_max_hops,
-            label_node_hops=label_node_hops, label_edge_hops=label_edge_hops,
-            label_seeds=label_seeds, to_fixed_point=to_fixed_point,
-            direction=direction, edge_match=edge_match,
-            source_node_match=source_node_match, destination_node_match=destination_node_match,
-            source_node_query=source_node_query, destination_node_query=destination_node_query,
-            edge_query=edge_query, return_as_wave_front=return_as_wave_front,
-            include_zero_hop_seed=include_zero_hop_seed, target_wave_front=target_wave_front,
-        )
+        from graphistry.compute.gfql.engine_polars.gpu import gpu_mode
+        with gpu_mode(engine_concrete == Engine.POLARS_GPU):
+            return hop_polars(
+                self, nodes, hops,
+                min_hops=min_hops, max_hops=max_hops,
+                output_min_hops=output_min_hops, output_max_hops=output_max_hops,
+                label_node_hops=label_node_hops, label_edge_hops=label_edge_hops,
+                label_seeds=label_seeds, to_fixed_point=to_fixed_point,
+                direction=direction, edge_match=edge_match,
+                source_node_match=source_node_match, destination_node_match=destination_node_match,
+                source_node_query=source_node_query, destination_node_query=destination_node_query,
+                edge_query=edge_query, return_as_wave_front=return_as_wave_front,
+                include_zero_hop_seed=include_zero_hop_seed, target_wave_front=target_wave_front,
+            )
 
     def _combine_first_no_warn(target, fill):
         """Avoid pandas concat warning when combine_first sees empty inputs."""