diff --git a/graphistry/Engine.py b/graphistry/Engine.py index 3209f5349c..19bcfdcef4 100644 --- a/graphistry/Engine.py +++ b/graphistry/Engine.py @@ -15,6 +15,8 @@ # triggers ``graphistry.compute`` package init at runtime (would be circular). from graphistry.compute.typing import DataFrameT, PolarsFrame, PolarsSeriesT, SeriesT from graphistry.Plottable import Plottable + import cudf + import cupy import polars as pl # TypeGuard, NOT TypeIs (PEP 742). TypeIs additionally narrows the negative branch, and to # do that soundly it REQUIRES the narrowed type to be consistent with the declared input @@ -249,6 +251,32 @@ def is_polars_df(df: object) -> "TypeGuard[PolarsFrame]": return df is not None and "polars" in type(df).__module__ +def is_cudf_series(series: object) -> "TypeGuard[cudf.Series]": + """True if ``series`` is a cuDF ``Series``. + + The module check avoids importing the optional cuDF dependency for other engines; + the exact ``isinstance`` check prevents look-alike classes from taking the cuDF path. + """ + if series is None or "cudf" not in type(series).__module__: + return False + import cudf + + return isinstance(series, cudf.Series) + + +def is_cupy_array(array: object) -> "TypeGuard[cupy.ndarray]": + """True if ``array`` is a CuPy array. + + The module check avoids importing the optional GPU dependency for other + engines; the exact ``isinstance`` check rejects look-alikes. + """ + if array is None or "cupy" not in type(array).__module__: + return False + import cupy + + return isinstance(array, cupy.ndarray) + + def active_frames_are_polars(g: Any) -> bool: """True if ``g``'s active table (nodes, else edges) is a polars frame.""" if g._nodes is not None: diff --git a/graphistry/compute/gfql/index/api.py b/graphistry/compute/gfql/index/api.py index 8fa76166bd..aa39d51e0f 100644 --- a/graphistry/compute/gfql/index/api.py +++ b/graphistry/compute/gfql/index/api.py @@ -8,12 +8,16 @@ from __future__ import annotations import copy -from typing import Dict, List, Literal, Mapping, Optional, Sequence, Set, Tuple, cast +from typing import TYPE_CHECKING, Dict, List, Literal, Mapping, Optional, Sequence, Set, Tuple, Union, cast +import numpy as np import pandas as pd -from graphistry.Engine import EngineAbstract, Engine, EngineAbstractType, POLARS_ENGINES, is_polars_df, resolve_engine -from graphistry.compute.typing import DataFrameT +from graphistry.Engine import ( + EngineAbstract, Engine, EngineAbstractType, POLARS_ENGINES, + is_cudf_series, is_cupy_array, is_polars_df, is_polars_series, resolve_engine, +) +from graphistry.compute.typing import DataFrameT, SeriesT from graphistry.Plottable import Plottable from .registry import ( AdjacencyIndex, ColStatsFact, PartitionValue, ColStatsRole, GfqlIndexRegistry, EMPTY_REGISTRY, NodeIdIndex, @@ -34,11 +38,31 @@ TraceEngine, SidecarFactKind, ) +if TYPE_CHECKING: + import cudf + from graphistry.compute.typing import ArrayLike, PolarsSeriesT + + CrossEngineByteSourceT = Union[ArrayLike, SeriesT, PolarsSeriesT, "cudf.Series"] + # Private Plottable attachment keys. Keep access behind helpers. POLICY_ATTR = "_gfql_index_policy" REGISTRY_ATTR = "_gfql_index_registry" +def _byte_source_nbytes(value: CrossEngineByteSourceT) -> int: + if is_polars_series(value): + return int(value.estimated_size()) + if is_cudf_series(value): + return int(value.memory_usage()) + if isinstance(value, np.ndarray): + return int(value.nbytes) + if is_cupy_array(value): + return int(value.nbytes) + if isinstance(value, pd.Series): + return int(value.memory_usage()) + return 0 + + # --- lightweight, thread-local index decision trace (for gfql_explain) ------- import threading as _threading _TRACE = _threading.local() @@ -591,7 +615,7 @@ def show_indexes( ) usable, reason = _index_usability(kind, typed.engine, valid, query_engine) # other_values is shared with the ordinary adjacency index, not added memory. - nbytes = sum(int(getattr(typed, name).nbytes) for name in ( + nbytes = sum(_byte_source_nbytes(getattr(typed, name)) for name in ( "keys_sorted", "group_offsets", "row_positions", )) rows.append({ @@ -611,7 +635,7 @@ def show_indexes( is segment ) usable, reason = _index_usability(NODE_TYPE_SEGMENT, segment.engine, valid, query_engine) - nbytes = sum(int(getattr(segment, name).nbytes) for name in ( + nbytes = sum(_byte_source_nbytes(getattr(segment, name)) for name in ( "keys_sorted", "group_offsets", "row_positions", )) rows.append({ diff --git a/graphistry/tests/compute/gfql/index/test_node_type_segments.py b/graphistry/tests/compute/gfql/index/test_node_type_segments.py index 511bf8d044..e1fecf746f 100644 --- a/graphistry/tests/compute/gfql/index/test_node_type_segments.py +++ b/graphistry/tests/compute/gfql/index/test_node_type_segments.py @@ -1,8 +1,11 @@ """Node type-segment index must preserve canonical GFQL results (edge typed-index dual).""" from __future__ import annotations +import numpy as np import pandas as pd import pytest +import sys +import types import graphistry from graphistry.Engine import Engine @@ -10,10 +13,28 @@ from graphistry.compute.gfql.index import ( get_registry, index_trace, maybe_index_hop, with_index_policy, ) +from graphistry.Engine import is_cupy_array +from graphistry.compute.gfql.index.api import _byte_source_nbytes from graphistry.compute.gfql.index.cypher_ddl import parse_index_ddl from graphistry.compute.gfql.index.wire import apply_index_op +class _CpuCudfSeries: + def __init__(self, values: list[int]) -> None: + self.values = values + + def memory_usage(self) -> int: + return 16 + + +class _CpuCuPyArray: + nbytes = 8 + + +_CpuCudfSeries.__module__ = "cudf" +_CpuCuPyArray.__module__ = "cupy" + + def _graph(engine: str): # Node 1 has TWO labels (membership); node 5 has a null type; node 4 is # label-only-False; a null type never matches any scalar segment. @@ -37,6 +58,79 @@ def _graph(engine: str): return graphistry.bind(node="id", source="src", destination="dst").nodes(nodes).edges(edges) +def test_node_segment_show_indexes_reports_cudf_bytes() -> None: + pytest.importorskip("cudf") + cupy = pytest.importorskip("cupy") + try: + device_count = cupy.cuda.runtime.getDeviceCount() + except Exception: + pytest.skip("CUDA device detection failed") + if device_count == 0: + pytest.skip("cuDF requires a CUDA device") + indexed = _indexed(_graph("cudf"), "cudf") + rows = indexed.show_indexes() + segment_rows = rows[rows["kind"] == "node_type_segment"] + assert not segment_rows.empty + assert segment_rows["nbytes"].sum() > 0 + + +def test_byte_source_nbytes_uses_pandas_measurement() -> None: + series = pd.Series([1], dtype="int64") + assert _byte_source_nbytes(series) == series.memory_usage() + + +def test_byte_source_nbytes_uses_polars_measurement() -> None: + pl = pytest.importorskip("polars") + series = pl.Series([1], dtype=pl.Int64) + assert _byte_source_nbytes(series) == series.estimated_size() + + +def test_byte_source_nbytes_uses_cudf_measurement() -> None: + cudf = pytest.importorskip("cudf") + try: + device_count = cudf.cuda.runtime.getDeviceCount() + except Exception: + pytest.skip("CUDA device detection failed") + if device_count == 0: + pytest.skip("cuDF requires a CUDA device") + series = cudf.Series([1], dtype="int64") + assert _byte_source_nbytes(series) == series.memory_usage() + + +def test_byte_source_nbytes_uses_cudf_guard_without_cuda(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setitem(sys.modules, "cudf", types.SimpleNamespace(Series=_CpuCudfSeries)) + series = _CpuCudfSeries([1]) + assert _byte_source_nbytes(series) == 16 + + +def test_byte_source_nbytes_defaults_to_zero_for_unknown_byte_source() -> None: + assert _byte_source_nbytes(object()) == 0 + + +def test_byte_source_nbytes_uses_numpy_measurement() -> None: + assert _byte_source_nbytes(np.asarray([1], dtype="int64")) == 8 + + +def test_byte_source_nbytes_uses_cupy_measurement() -> None: + cupy = pytest.importorskip("cupy") + try: + device_count = cupy.cuda.runtime.getDeviceCount() + except Exception: + pytest.skip("CUDA device detection failed") + if device_count == 0: + pytest.skip("CuPy requires a CUDA device") + array = cupy.asarray([1], dtype="int64") + assert _byte_source_nbytes(array) == 8 + assert is_cupy_array(array) + + +def test_byte_source_nbytes_uses_cupy_guard_without_cuda(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setitem(sys.modules, "cupy", types.SimpleNamespace(ndarray=_CpuCuPyArray)) + array = _CpuCuPyArray() + assert _byte_source_nbytes(array) == 8 + assert is_cupy_array(array) + + def _pandas(frame): return frame.to_pandas() if hasattr(frame, "to_pandas") else frame