diff --git a/.flake8 b/.flake8 index 959e9a22c..9ba580871 100644 --- a/.flake8 +++ b/.flake8 @@ -10,6 +10,7 @@ extend-ignore = SIM401, # E203 whitespace before ':' (to be compatible with black) per-file-ignores = + benchmarks/*.py:B015,B018, scripts/create_pickle.py:F403,F405, graphblas/tests/*.py:T201,B043, graphblas/core/ss/matrix.py:SIM113, diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 3582142d2..b9603a6b7 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -26,7 +26,11 @@ repos: - id: check-illegal-windows-names - id: check-merge-conflict - id: check-ast - - id: check-json # no JSON files yet; enabled for the future + # `asv.conf.json` is JSONC: asv documents `//` comments and strips them + # in its own loader, but this hook is a strict JSON parser. Excluding + # that one path keeps the hook live for every other JSON file. + - id: check-json + exclude: ^asv\.conf\.json$ - id: check-toml - id: check-yaml - id: check-executables-have-shebangs diff --git a/README.md b/README.md index 7b157c979..c9ab73d8f 100644 --- a/README.md +++ b/README.md @@ -178,7 +178,7 @@ s(accum) << v.reduce(op) ## Creating new Vectors / Matrices ```python -A = Matrix.new(dtype, num_rows, num_cols) # new_type +A = Matrix(dtype, num_rows, num_cols) # new_type B = A.dup() # dup A = Matrix.from_coo([row_indices], [col_indices], [values]) # build ``` @@ -225,7 +225,7 @@ Python-graphblas requires `numba` which enables compiling user-defined Python fu Example customized UnaryOp: ```python -from graphblas import unary +from graphblas import unary, Vector def force_odd_func(x): if x % 2 == 0: diff --git a/asv.conf.json b/asv.conf.json new file mode 100644 index 000000000..acd395737 --- /dev/null +++ b/asv.conf.json @@ -0,0 +1,73 @@ +{ + // asv (airspeed velocity) config for python-graphblas. + // + // Placement: this file and the benchmarks/ directory are intended to live at the + // repo root (next to pyproject.toml). The paths below are relative to that root. + // asv's config loader strips these // line comments; a strict JSON linter will + // not. If the repo's pre-commit `check-json` hook is enabled, exclude this file + // (asv.conf.json is asv's own format, and its docs use // comments). + // + // Typical use: + // asv run # benchmark current commit + // asv continuous main HEAD # compare a branch against main + // asv publish && asv preview # build/serve the HTML report + + "version": 1, + "project": "python-graphblas", + "project_url": "https://github.com/python-graphblas/python-graphblas", + + // "." means the repo that contains this file. python-graphblas is pure Python, so + // building each commit is a cheap editable install; the C library comes from the + // separate suitesparse-graphblas package, pinned in the matrix below. + "repo": ".", + "branches": ["main"], + "dvcs": "git", + + // conda matches the project's primary dev path and lets us pin the C library + // cleanly from conda-forge. asv will use mamba automatically if it is installed. + "environment_type": "conda", + "conda_channels": ["conda-forge"], + + // PINNING NOTE: python-suitesparse-graphblas (the C library) is pinned so a run + // measures python-graphblas-side regressions in isolation. To instead track + // end-to-end performance including C-library changes, drop the pin (set it to []) + // or bump it deliberately. Keep this in sync with the PSG version the suite is + // calibrated against (see scripts/ci_pick_versions.py PSG pools). + "matrix": { + "req": { + "python-suitesparse-graphblas": ["10.3.1.0"], + "numpy": [], + "scipy": [], + "pandas": [], + "numba": [] + } + }, + + // python-graphblas is pure Python: a no-build-isolation editable install is fast + // and reuses the already-solved conda env (which supplies the C library plus + // numba/scipy/pandas). asv substitutes {build_dir} for the checkout of each commit. + "build_command": [], + "install_command": [ + "in-dir={build_dir} python -m pip install --no-build-isolation --no-deps -e ." + ], + "uninstall_command": [ + "in-dir={build_dir} python -m pip uninstall -y python-graphblas" + ], + + "benchmark_dir": "benchmarks", + "env_dir": ".asv/env", + "results_dir": ".asv/results", + "html_dir": ".asv/html", + + // Give the C-library install room to solve/download. + "install_timeout": 900, + + // Isolate JIT/compiler and thread-count noise so numbers are comparable across + // runs. GraphBLAS is multithreaded by default; pin to 1 thread for reproducible + // single-core timings (raise this in a dedicated run if you want to benchmark + // parallel scaling instead). + "environment_variables": { + "OMP_NUM_THREADS": "1", + "GRAPHBLAS_TEST_SEED": "0" + } +} diff --git a/benchmarks/README.md b/benchmarks/README.md new file mode 100644 index 000000000..5962927f7 --- /dev/null +++ b/benchmarks/README.md @@ -0,0 +1,91 @@ +# asv benchmark scaffold for python-graphblas + +A starting airspeed velocity (asv) benchmark suite covering the library's hot +paths. Built as a scaffold: review, then move `asv.conf.json` and `benchmarks/` +to the repo root (next to `pyproject.toml`). `_verify.py` is a local sanity +checker and does not need to ship. + +## Layout + +``` +asv.conf.json asv config, tuned for this repo (see caveats below). + JSONC: asv documents `//` comments and strips them in + its loader, so `.pre-commit-config.yaml` excludes this + one path from the strict `check-json` hook. asv also + accepts `asv.conf.jsonc`, which would need no exclude, + but asv is not a pinned dependency here and older + versions resolve only the `.json` name. +benchmarks/ + __init__.py package marker + note on the dual-import shim + common.py shared, seeded data builders (called from setup, not timed) + scalar_access.py single-element get / extract / assign, hit vs miss + index_parse.py parse_index int fast lane vs numpy-int lane (index build) + small_ops.py ewise / apply / reduce on ~100-element objects (overhead) + large_kernels.py mxm / mxv / ewise / reduce on ~1e6 nnz (C-library bound) + conversions.py from_coo/to_coo, from_dense/to_dense, scipy interop + imports.py cold import + first-use timing (timeraw, fresh subprocess) + repr_bench.py repr / _repr_html_ for small and large objects +_verify.py standalone check: runs every benchmark once (not an asv file) +``` + +## Running + +```bash +pip install asv +asv machine --yes # one-time: record machine info +asv run # benchmark the current commit +asv continuous main HEAD # compare a branch against main, flag regressions +asv publish && asv preview # build and serve the HTML report +asv run --bench scalar_access # run a subset by regex +``` + +Quick correctness check without asv (runs each benchmark once): + +```bash +python _verify.py +``` + +## Design choices + +- `setup()` builds all data once, outside the timed region; every builder is + seeded so runs are comparable across commits. +- Sizes: "small" is 10 to 1000 elements (overhead-dominated, so we track + Python-side cost); "large" is ~1e6 nonzeros (C-library dominated). Average + matrix degree is ~1 so `mxm` output stays near the input size instead of + exploding into quadratic fill-in. Verified per-call times on one dev machine: + large kernels ran in 1 to 5 ms, conversions under 10 ms, cold `import +graphblas` ~32 ms, init/first-operator ~0.2 s. Bump `common.LARGE_NNZ` or add a + higher-degree matrix if you want heavier `mxm`. +- Large kernels use `number = 1`, `warmup_time = 0`, small `repeat`, and a + `timeout`, so asv does not batch them into multi-second runs. +- `imports.py` uses `timeraw_*`, which runs the returned code string in a fresh + subprocess. That is the only way to measure true cold import cost; a normal + benchmark would read ~0 because the module is already imported. +- `asv.conf.json` pins `python-suitesparse-graphblas` so a run isolates + python-graphblas-side regressions from C-library changes. Drop the pin to track + end-to-end (library + C) performance. Keep the pinned version in sync with the + PSG pools in `scripts/ci_pick_versions.py`. +- `OMP_NUM_THREADS=1` pins GraphBLAS to one thread for reproducible single-core + timings. Raise it in a dedicated run to benchmark parallel scaling. + +## Caveats to resolve before merging + +- **dev deps**: `asv` is deliberately NOT added to `dev-requirements.txt` or + `environment.yml`. Nothing in the test suite or CI invokes it, so adding it + would make every contributor's dev install fetch a package only used by + people who opt into benchmarking; and this repo ties dependency changes to + `scripts/check_versions.sh` and the CI version pools, which a benchmarks-only + change has no business editing. `pip install asv` is in the usage section + above. Adding a `benchmark` optional-dependencies group (so that + `pip install python-graphblas[benchmark]` works) is the natural move if the + suite graduates from scaffold to something CI runs. +- **CI wiring**: asv is not wired into CI here. A weekly cron running `asv +continuous` against a fixed machine (or asv's own regression detection) is the + natural follow-up; per-PR benchmarking on shared GitHub runners is too noisy to + gate on. +- **First-run JIT/compile noise**: the first use of some operators triggers + numba/JIT work. `setup()` touches the operators the benchmark uses, but the + very first `asv run` on a fresh env may still show inflated one-off numbers; + asv's repeats and its own warmup mitigate this. +- **pandas dependency**: `repr_bench.py` needs pandas (the default/test extras + include it). The env matrix installs pandas, so this is satisfied. diff --git a/benchmarks/__init__.py b/benchmarks/__init__.py new file mode 100644 index 000000000..dd8e37ffb --- /dev/null +++ b/benchmarks/__init__.py @@ -0,0 +1,12 @@ +"""python-graphblas benchmark suite (airspeed velocity). + +Benchmark modules live alongside this file. Shared data builders are in +``common.py``. Each benchmark module imports ``common`` with a dual path so it +works whether asv imports the files as a package or adds ``benchmark_dir`` to +``sys.path`` and imports them flat: + + try: + from . import common + except ImportError: # imported flat, not as a package + import common +""" diff --git a/benchmarks/_verify.py b/benchmarks/_verify.py new file mode 100755 index 000000000..cb9ccd5d4 --- /dev/null +++ b/benchmarks/_verify.py @@ -0,0 +1,107 @@ +#!/usr/bin/env python +"""Standalone sanity check for the asv benchmark suite (not an asv benchmark). + +Imports every benchmark module the way asv does (benchmark_dir on sys.path), +then for each benchmark class runs setup() and calls each benchmark method once +across all parameter combinations. timeraw_* methods return a code string, which +is executed in a fresh subprocess. Prints PASS/FAIL per benchmark and a timing so +we can spot anything accidentally slow. + +Usage: python _verify.py +""" + +import importlib +import inspect +import itertools +import subprocess +import sys +import time +from pathlib import Path + +HERE = Path(__file__).resolve().parent +BENCH_DIR = HERE / "benchmarks" +sys.path.insert(0, str(BENCH_DIR)) # flat import mode, as asv does + +MODULES = [ + "scalar_access", + "index_parse", + "small_ops", + "large_kernels", + "conversions", + "imports", + "repr_bench", +] +PREFIXES = ("time_", "timeraw_", "peakmem_", "mem_", "track_") + + +def param_combos(cls): + params = getattr(cls, "params", None) + if not params: + return [()] + # asv: a flat list is a single parameter; a list of lists is a product. + if params and isinstance(params[0], (list, tuple)): + return list(itertools.product(*params)) + return [(p,) for p in params] + + +def run_one(cls, combo): + results = [] + obj = cls() + if hasattr(obj, "setup"): + obj.setup(*combo) + for name in sorted(dir(obj)): + if not name.startswith(PREFIXES): + continue + method = getattr(obj, name) + if not callable(method): + continue + label = f"{cls.__module__}.{cls.__name__}.{name}{combo or ''}" + t0 = time.perf_counter() + try: + if name.startswith("timeraw_"): + code = method() + setup_code = "" + if isinstance(code, tuple): + code, setup_code = code + script = (setup_code + "\n" + code) if setup_code else code + subprocess.run( + [sys.executable, "-c", script], check=True, capture_output=True, timeout=180 + ) + else: + method(*combo) + dt = time.perf_counter() - t0 + results.append((label, dt, None)) + except Exception as exc: + dt = time.perf_counter() - t0 + results.append((label, dt, repr(exc))) + if hasattr(obj, "teardown"): + obj.teardown(*combo) + return results + + +def main(): + all_results = [] + for modname in MODULES: + mod = importlib.import_module(modname) + classes = [ + c for _, c in inspect.getmembers(mod, inspect.isclass) if c.__module__ == modname + ] + for cls in classes: + for combo in param_combos(cls): + all_results.extend(run_one(cls, combo)) + + fails = [r for r in all_results if r[2] is not None] + for label, dt, err in all_results: + status = "PASS" if err is None else "FAIL" + line = f"[{status}] {dt:7.3f}s {label}" + if err is not None: + line += f" -> {err}" + print(line) + print( + f"\n{len(all_results) - len(fails)}/{len(all_results)} benchmarks ran; {len(fails)} failed" + ) + return 1 if fails else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/benchmarks/common.py b/benchmarks/common.py new file mode 100644 index 000000000..44db815fa --- /dev/null +++ b/benchmarks/common.py @@ -0,0 +1,91 @@ +"""Shared data builders for the python-graphblas asv benchmark suite. + +Each benchmark module imports these helpers and calls them from ``setup`` so the +input data is constructed once per benchmark (outside the timed region). All +builders take a seed and use ``numpy.random.default_rng`` so runs are +deterministic and comparable across commits. + +Size constants are chosen so the "large" kernels land near ~1e6 nonzeros (enough +to exercise the C library) while still finishing in well under a second each, and +so no dense intermediate blows up memory (a 1e6 x 1e6 dense matrix is never +materialized; dense conversions use a separate, small, fully dense matrix). +""" + +import numpy as np + +import graphblas as gb +from graphblas import Matrix, Vector + +# Sizes + +# "Large" sparse operands: square matrix and vector near 1e6 nonzeros. Average +# degree ~1 keeps mxm output bounded (A @ A stays near 1e6 nnz) so the kernel +# benchmarks do not accidentally measure quadratic fill-in. +LARGE_N = 1_000_000 +LARGE_NNZ = 1_000_000 + +# "Small" operands: overhead-dominated. At this size the wall time is almost all +# Python-side expression machinery, which is exactly what we want to track. +SMALL_SIZE = 100 + +# Dense conversions use a fully dense square matrix small enough to materialize. +DENSE_DIM = 1000 # 1e6 dense elements + +# scipy interop matrix: ~1e6 nnz at 1% density. +SCIPY_DIM = 10_000 +SCIPY_DENSITY = 0.01 + + +# Builders + + +def make_coo(n, nnz, seed=0, dtype="FP64"): + """Return (rows, cols, vals) numpy arrays for an n x n matrix with ~nnz entries. + + Duplicates are possible; callers that build a Matrix should pass a ``dup_op``. + """ + rng = np.random.default_rng(seed) + rows = rng.integers(0, n, nnz, dtype=np.uint64) + cols = rng.integers(0, n, nnz, dtype=np.uint64) + if dtype == "BOOL": + vals = rng.integers(0, 2, nnz, dtype=np.bool_) + elif dtype in ("INT64", "INT32"): + vals = rng.integers(1, 100, nnz, dtype=np.int64) + else: + vals = rng.random(nnz) + return rows, cols, vals + + +def make_matrix(n=LARGE_N, nnz=LARGE_NNZ, seed=0, dtype="FP64"): + """Square Matrix with ~nnz entries (duplicates summed).""" + rows, cols, vals = make_coo(n, nnz, seed=seed, dtype=dtype) + dup = gb.binary.lor if dtype == "BOOL" else gb.binary.plus + return Matrix.from_coo(rows, cols, vals, nrows=n, ncols=n, dtype=dtype, dup_op=dup) + + +def make_vector(size=LARGE_N, nnz=LARGE_NNZ, seed=1, dtype="FP64"): + """Vector of length ``size`` with ~nnz entries (duplicates summed).""" + rng = np.random.default_rng(seed) + idx = rng.integers(0, size, nnz, dtype=np.uint64) + vals = rng.random(nnz) if dtype not in ("INT64", "INT32") else rng.integers(1, 100, nnz) + dup = gb.binary.plus + return Vector.from_coo(idx, vals, size=size, dtype=dtype, dup_op=dup) + + +def make_dense_vector(size, seed=2, dtype="FP64"): + """Fully dense Vector (no missing entries), e.g. an mxv operand.""" + rng = np.random.default_rng(seed) + return Vector.from_dense(rng.random(size)) + + +def make_dense_matrix(dim=DENSE_DIM, seed=3): + """Fully dense square Matrix built from a numpy array.""" + rng = np.random.default_rng(seed) + return Matrix.from_dense(rng.random((dim, dim))) + + +def make_scipy(dim=SCIPY_DIM, density=SCIPY_DENSITY, seed=4, fmt="csr"): + """A scipy.sparse matrix for interop benchmarks.""" + import scipy.sparse as sp + + return sp.random(dim, dim, density=density, format=fmt, random_state=seed) diff --git a/benchmarks/conversions.py b/benchmarks/conversions.py new file mode 100644 index 000000000..90be6c986 --- /dev/null +++ b/benchmarks/conversions.py @@ -0,0 +1,98 @@ +"""Import/export conversions between graphblas objects and numpy/scipy formats. + +Sparse conversions (from_coo/to_coo, scipy interop) run on ~1e6 nonzeros. Dense +conversions use a separate fully dense 1000x1000 matrix so nothing materializes a +1e6 x 1e6 dense array. +""" + +import numpy as np + +from graphblas import Matrix, Vector, io + +try: + from . import common +except ImportError: + import common + + +class MatrixCoo: + number = 1 + repeat = 5 + warmup_time = 0 + timeout = 300 + + def setup(self): + self.rows, self.cols, self.vals = common.make_coo(common.LARGE_N, common.LARGE_NNZ, seed=40) + self.M = common.make_matrix(seed=40) + + def time_from_coo(self): + Matrix.from_coo( + self.rows, + self.cols, + self.vals, + nrows=common.LARGE_N, + ncols=common.LARGE_N, + dup_op="plus", + ) + + def time_to_coo(self): + self.M.to_coo() + + +class MatrixDense: + number = 1 + repeat = 5 + warmup_time = 0 + timeout = 300 + + def setup(self): + self.dense = np.random.default_rng(41).random((common.DENSE_DIM, common.DENSE_DIM)) + self.M = Matrix.from_dense(self.dense) + + def time_from_dense(self): + Matrix.from_dense(self.dense) + + def time_to_dense(self): + self.M.to_dense() + + +class MatrixScipy: + number = 1 + repeat = 5 + warmup_time = 0 + timeout = 300 + + def setup(self): + self.sp = common.make_scipy(seed=42) + self.M = io.from_scipy_sparse(self.sp) + + def time_from_scipy_sparse(self): + io.from_scipy_sparse(self.sp) + + def time_to_scipy_sparse(self): + io.to_scipy_sparse(self.M, format="csr") + + +class VectorConvert: + number = 1 + repeat = 5 + warmup_time = 0 + timeout = 300 + + def setup(self): + self.v = common.make_vector(seed=43) + self.idx, self.vals = self.v.to_coo() + self.dense_arr = np.random.default_rng(44).random(common.LARGE_N) + self.dv = Vector.from_dense(self.dense_arr) + + def time_from_coo(self): + Vector.from_coo(self.idx, self.vals, size=common.LARGE_N, dup_op="plus") + + def time_to_coo(self): + self.v.to_coo() + + def time_from_dense(self): + Vector.from_dense(self.dense_arr) + + def time_to_dense(self): + self.dv.to_dense() diff --git a/benchmarks/imports.py b/benchmarks/imports.py new file mode 100644 index 000000000..9204b26a2 --- /dev/null +++ b/benchmarks/imports.py @@ -0,0 +1,35 @@ +"""Import and cold-start timing. + +These use asv's ``timeraw_*`` form: each returns a code string that asv runs in a +*fresh* subprocess, so the measurement is the true cold cost (module already +imported into the benchmark process would otherwise read as ~0). ``timeraw`` +benchmarks cannot see ``setup`` state or this module's imports, so everything +they need must be inside the returned string. + +``graphblas`` initializes lazily: ``import graphblas`` is cheap, and the real +cost (loading the C library, building operator namespaces) is deferred until +first use. The staged benchmarks below separate those phases. +""" + + +class ImportTiming: + # A little headroom: importing numba/llvmlite on first operator use is not fast. + timeout = 120 + + def timeraw_import_graphblas(self): + return "import graphblas" + + def timeraw_from_import_core_types(self): + return "from graphblas import Matrix, Vector, Scalar" + + def timeraw_import_then_init(self): + # Force backend initialization (loads the C library). + return "import graphblas as gb; gb.init('suitesparse')" + + def timeraw_first_operator_access(self): + # First attribute access on an operator namespace triggers its lazy build. + return "import graphblas as gb; gb.binary.plus" + + def timeraw_first_matrix(self): + # End-to-end cold path: import, init, and build one tiny Matrix. + return "import graphblas as gb; gb.Matrix.from_coo([0], [0], [1.0], nrows=1, ncols=1)" diff --git a/benchmarks/index_parse.py b/benchmarks/index_parse.py new file mode 100644 index 000000000..5b5b0c856 --- /dev/null +++ b/benchmarks/index_parse.py @@ -0,0 +1,90 @@ +"""Index parsing overhead: the ``parse_index`` int fast lane. + +Every ``v[i]`` / ``A[i, j]`` builds an ``IndexerResolver`` that runs each index +through ``parse_index``. A plain Python ``int`` takes a dedicated fast lane that +skips two ``np.issubdtype`` checks (~110ns each) used by the numpy-integer lane. +These benchmarks isolate that parsing cost from element extraction (no ``.new()``, +no value read), so a regression in the fast lane, or an index accidentally +falling out of it, is visible on its own. + +Two granularities are measured: + +* the public path (``v[i]`` returns an index expression without resolving it), and +* ``IndexerResolver(obj, idx)`` directly, the tightest view of ``parse_index``. + +The plain-int and numpy-int variants sit side by side so the fast lane's margin +over the general lane is tracked directly. +""" + +import numpy as np + +from graphblas.core.expr import IndexerResolver + +try: + from . import common +except ImportError: # imported flat by asv, not as a package + import common + + +class VectorIndexParse: + """Parse a single vector index: plain-int fast lane vs numpy-int lane.""" + + def setup(self): + self.v = common.make_vector(size=10_000, nnz=1_000, seed=12) + self.i = 4321 # in range, positive + self.neg = -1 # in range, triggers the negative-wrap branch + self.npi = np.int64(4321) + for _ in range(3): + self.v[self.i] + self.v[self.neg] + self.v[self.npi] + IndexerResolver(self.v, self.i) + IndexerResolver(self.v, self.npi) + + # Public path: build the index expression (parse_index + expression object). + def time_getitem_int(self): + self.v[self.i] + + def time_getitem_int_negative(self): + self.v[self.neg] + + def time_getitem_numpy_int(self): + self.v[self.npi] + + # Tightest view: just the resolver (parse_index, no expression object). + def time_resolver_int(self): + IndexerResolver(self.v, self.i) + + def time_resolver_numpy_int(self): + IndexerResolver(self.v, self.npi) + + +class MatrixIndexParse: + """Parse a two-axis matrix index: plain-int fast lane vs numpy-int lane.""" + + def setup(self): + self.M = common.make_matrix(n=10_000, nnz=1_000, seed=13) + self.ij = (4321, 8765) + self.neg = (-1, -1) + self.npij = (np.int64(4321), np.int64(8765)) + for _ in range(3): + self.M[self.ij[0], self.ij[1]] + self.M[self.neg[0], self.neg[1]] + self.M[self.npij[0], self.npij[1]] + IndexerResolver(self.M, self.ij) + IndexerResolver(self.M, self.npij) + + def time_getitem_int(self): + self.M[self.ij[0], self.ij[1]] + + def time_getitem_int_negative(self): + self.M[self.neg[0], self.neg[1]] + + def time_getitem_numpy_int(self): + self.M[self.npij[0], self.npij[1]] + + def time_resolver_int(self): + IndexerResolver(self.M, self.ij) + + def time_resolver_numpy_int(self): + IndexerResolver(self.M, self.npij) diff --git a/benchmarks/large_kernels.py b/benchmarks/large_kernels.py new file mode 100644 index 000000000..f0c278409 --- /dev/null +++ b/benchmarks/large_kernels.py @@ -0,0 +1,73 @@ +"""Large kernels (~1e6 nonzeros): dominated by the C library, so these catch +suitesparse-graphblas-side regressions (and any python-graphblas overhead that +scales with data). + +Each op runs once per sample (``number = 1``) with no warmup, because the inputs +are large enough that a single call is well above timer resolution and we do not +want asv auto-tuning ``number`` up into multi-second batches. ``repeat`` and the +timeout keep total time bounded. Average degree ~1 keeps ``mxm`` output near the +input size instead of exploding into quadratic fill-in. +""" + +from graphblas import binary, monoid, semiring, unary + +try: + from . import common +except ImportError: + import common + + +class VectorLarge: + number = 1 + repeat = 5 + warmup_time = 0 + timeout = 300 + + def setup(self): + self.v = common.make_vector(seed=30) + self.u = common.make_vector(seed=31) + + def time_ewise_mult(self): + self.v.ewise_mult(self.u, binary.times).new() + + def time_ewise_add(self): + self.v.ewise_add(self.u, monoid.plus).new() + + def time_apply(self): + self.v.apply(unary.abs).new() + + def time_reduce(self): + self.v.reduce(monoid.plus).new() + + def peakmem_ewise_add(self): + self.v.ewise_add(self.u, monoid.plus).new() + + +class MatrixLarge: + number = 1 + repeat = 5 + warmup_time = 0 + timeout = 300 + + def setup(self): + self.A = common.make_matrix(seed=32) + self.B = common.make_matrix(seed=33) + self.x = common.make_dense_vector(common.LARGE_N, seed=34) + + def time_mxv(self): + self.A.mxv(self.x, semiring.plus_times).new() + + def time_mxm(self): + self.A.mxm(self.B, semiring.plus_times).new() + + def time_ewise_add(self): + self.A.ewise_add(self.B, monoid.plus).new() + + def time_reduce_rowwise(self): + self.A.reduce_rowwise(monoid.plus).new() + + def time_reduce_scalar(self): + self.A.reduce_scalar(monoid.plus).new() + + def peakmem_mxm(self): + self.A.mxm(self.B, semiring.plus_times).new() diff --git a/benchmarks/repr_bench.py b/benchmarks/repr_bench.py new file mode 100644 index 000000000..4f138a354 --- /dev/null +++ b/benchmarks/repr_bench.py @@ -0,0 +1,51 @@ +"""repr / _repr_html_ rendering. + +Object display goes through pandas and, for large objects, the formatting code +that decides what to elide. Small reprs are an overhead microbenchmark; large +reprs check that the truncation path stays cheap (it must not render every one of +~1e6 nonzeros). +""" + +from graphblas import Matrix, Vector + +try: + from . import common +except ImportError: + import common + + +class ReprSmall: + def setup(self): + self.M = Matrix.from_coo([0, 1, 2], [0, 1, 2], [1.0, 2.0, 3.0], nrows=4, ncols=4) + self.v = Vector.from_coo([0, 2], [1.0, 2.0], size=5) + + def time_repr_matrix(self): + repr(self.M) + + def time_repr_vector(self): + repr(self.v) + + def time_repr_html_matrix(self): + self.M._repr_html_() + + def time_repr_html_vector(self): + self.v._repr_html_() + + +class ReprLarge: + # Should be bounded by truncation, but give it room in case a regression makes + # it render the whole object. + timeout = 120 + + def setup(self): + self.M = common.make_matrix(seed=50) + self.v = common.make_vector(seed=51) + + def time_repr_matrix(self): + repr(self.M) + + def time_repr_vector(self): + repr(self.v) + + def time_repr_html_matrix(self): + self.M._repr_html_() diff --git a/benchmarks/scalar_access.py b/benchmarks/scalar_access.py new file mode 100644 index 000000000..0de089802 --- /dev/null +++ b/benchmarks/scalar_access.py @@ -0,0 +1,168 @@ +"""Single-element access: the tightest hot loops in the library. + +Extracting or assigning one element goes through the full expression + descriptor +machinery, so these times are almost entirely Python overhead. "hit" indexes an +element that is present; "miss" indexes an empty slot (the extract returns an +empty Scalar / ``None`` value). + +Several of these paths gained dedicated fast lanes (single-element extract, +``get``, ``__contains__``, integer ``__setitem__``); the benchmarks below track +each one so a regression that re-routes through the slow expression path shows up. +Every ``setup`` warms the path once so the first timed sample is not inflated by +one-off numba/dtype/operator cache population. +""" + +import graphblas as gb +from graphblas import Scalar + +try: + from . import common +except ImportError: # imported flat by asv, not as a package + import common + + +class ScalarObject: + """Scalar construction and .value round-trips.""" + + def setup(self): + self.s = Scalar.from_value(3.14, dtype=gb.dtypes.FP64) + # Warm construction / value round-trip caches. + for _ in range(3): + Scalar.from_value(3.14, dtype=gb.dtypes.FP64) + self.s.value + self.s.value = 2.0 + self.s.dup() + + def time_from_value(self): + Scalar.from_value(3.14, dtype=gb.dtypes.FP64) + + def time_get_value(self): + self.s.value + + def time_set_value(self): + self.s.value = 2.0 + + def time_dup(self): + self.s.dup() + + +class VectorElement: + """Get / extract / assign a single vector element, present vs missing.""" + + def setup(self): + # size 10000, ~1000 present entries; index 0 forced present, 1 forced empty + self.v = common.make_vector(size=10_000, nnz=1_000, seed=10) + self.v[0] = 1.0 + del self.v[1] + self.hit = 0 + self.miss = 1 + # Warm each timed path once (extract, get, contains, setitem, the value / + # float / int single-extract fast paths) so no timed sample pays the + # first-call numba compile or cache-fill cost. + for _ in range(3): + self.v[self.hit].new() + self.v[self.miss].new() + self.v[self.hit].value + self.v[self.miss].value + float(self.v[self.hit]) + int(self.v[self.hit]) + self.v.get(self.hit) + self.v.get(self.miss) + self.v.get(self.miss, 0.0) + _ = self.hit in self.v + _ = self.miss in self.v + self.v[self.hit] = 5.0 + + def time_getitem_hit(self): + self.v[self.hit].new() + + def time_getitem_miss(self): + self.v[self.miss].new() + + def time_value_hit(self): + self.v[self.hit].value + + def time_value_miss(self): + self.v[self.miss].value + + def time_float_hit(self): + float(self.v[self.hit]) + + def time_int_hit(self): + int(self.v[self.hit]) + + def time_get_hit(self): + self.v.get(self.hit) + + def time_get_miss(self): + self.v.get(self.miss) + + def time_get_miss_default(self): + # `get` with an explicit default; the miss returns the default rather than None. + self.v.get(self.miss, 0.0) + + def time_contains_hit(self): + self.hit in self.v + + def time_contains_miss(self): + self.miss in self.v + + def time_setitem(self): + # Overwrites an existing slot, so state is stable across repeated calls. + self.v[self.hit] = 5.0 + + +class MatrixElement: + """Get / extract / assign a single matrix element, present vs missing.""" + + def setup(self): + self.M = common.make_matrix(n=10_000, nnz=1_000, seed=11) + self.M[0, 0] = 1.0 + del self.M[1, 1] + self.hit = (0, 0) + self.miss = (1, 1) + for _ in range(3): + self.M[self.hit[0], self.hit[1]].new() + self.M[self.miss[0], self.miss[1]].new() + self.M[self.hit[0], self.hit[1]].value + self.M[self.miss[0], self.miss[1]].value + float(self.M[self.hit[0], self.hit[1]]) + self.M.get(self.hit[0], self.hit[1]) + self.M.get(self.miss[0], self.miss[1]) + self.M.get(self.miss[0], self.miss[1], 0.0) + _ = self.hit in self.M + _ = self.miss in self.M + self.M[self.hit[0], self.hit[1]] = 5.0 + + def time_getitem_hit(self): + self.M[self.hit[0], self.hit[1]].new() + + def time_getitem_miss(self): + self.M[self.miss[0], self.miss[1]].new() + + def time_value_hit(self): + self.M[self.hit[0], self.hit[1]].value + + def time_value_miss(self): + self.M[self.miss[0], self.miss[1]].value + + def time_float_hit(self): + float(self.M[self.hit[0], self.hit[1]]) + + def time_get_hit(self): + self.M.get(self.hit[0], self.hit[1]) + + def time_get_miss(self): + self.M.get(self.miss[0], self.miss[1]) + + def time_get_miss_default(self): + self.M.get(self.miss[0], self.miss[1], 0.0) + + def time_contains_hit(self): + self.hit in self.M + + def time_contains_miss(self): + self.miss in self.M + + def time_setitem(self): + self.M[self.hit[0], self.hit[1]] = 5.0 diff --git a/benchmarks/small_ops.py b/benchmarks/small_ops.py new file mode 100644 index 000000000..ce2c040a6 --- /dev/null +++ b/benchmarks/small_ops.py @@ -0,0 +1,72 @@ +"""Small-operand op overhead. + +On ~100-element objects the C kernels are trivially fast, so these times isolate +the Python-side cost of building an expression and resolving operators. The +``time_build_*`` benchmarks stop at the expression object (no ``.new()``) to +separate expression construction from evaluation. +""" + +from graphblas import binary, monoid, semiring, unary + +try: + from . import common +except ImportError: + import common + + +class SmallVector: + params = [10, 100, 1000] + param_names = ["size"] + + def setup(self, size): + # Dense vectors so every op touches ``size`` elements. + self.v = common.make_dense_vector(size, seed=20) + self.u = common.make_dense_vector(size, seed=21) + + def time_build_ewise_mult(self, size): + # Expression only, not evaluated: pure construction overhead. + self.v.ewise_mult(self.u, binary.times) + + def time_ewise_mult(self, size): + self.v.ewise_mult(self.u, binary.times).new() + + def time_ewise_add(self, size): + self.v.ewise_add(self.u, monoid.plus).new() + + def time_apply(self, size): + self.v.apply(unary.abs).new() + + def time_apply_bind_scalar(self, size): + self.v.apply(binary.plus, right=1.0).new() + + def time_reduce(self, size): + self.v.reduce(monoid.plus).new() + + def time_assign_into(self, size): + # The `<<` update path: evaluate into an existing object with no mask/accum. + self.v << self.v.ewise_mult(self.u, binary.times) + + +class SmallMatrix: + params = [10, 100] + param_names = ["dim"] + + def setup(self, dim): + # Dense dim x dim (dim**2 entries): 100 or 10000 nonzeros. + self.A = common.make_dense_matrix(dim=dim, seed=22) + self.B = common.make_dense_matrix(dim=dim, seed=23) + + def time_ewise_mult(self, dim): + self.A.ewise_mult(self.B, binary.times).new() + + def time_apply(self, dim): + self.A.apply(unary.abs).new() + + def time_reduce_rowwise(self, dim): + self.A.reduce_rowwise(monoid.plus).new() + + def time_reduce_scalar(self, dim): + self.A.reduce_scalar(monoid.plus).new() + + def time_mxm(self, dim): + self.A.mxm(self.B, semiring.plus_times).new() diff --git a/docs/api_reference/index.rst b/docs/api_reference/index.rst index 84e7d65eb..219e19bae 100644 --- a/docs/api_reference/index.rst +++ b/docs/api_reference/index.rst @@ -8,6 +8,8 @@ API Reference :maxdepth: 2 collections + types operators io + utilities exceptions diff --git a/docs/api_reference/io.rst b/docs/api_reference/io.rst index 1cfc98516..9c8b8e2cb 100644 --- a/docs/api_reference/io.rst +++ b/docs/api_reference/io.rst @@ -73,3 +73,7 @@ Visualization ~~~~~~~~~~~~~ .. autofunction:: graphblas.viz.draw + +.. autofunction:: graphblas.viz.spy + +.. autofunction:: graphblas.viz.datashade diff --git a/docs/api_reference/operators.rst b/docs/api_reference/operators.rst index 8836bb638..6853831cd 100644 --- a/docs/api_reference/operators.rst +++ b/docs/api_reference/operators.rst @@ -31,8 +31,20 @@ IndexUnaryOp .. autoclass:: graphblas.core.operator.IndexUnaryOp() :members: +IndexBinaryOp +~~~~~~~~~~~~~ + +.. autoclass:: graphblas.core.operator.IndexBinaryOp() + :members: + SelectOp ~~~~~~~~ .. autoclass:: graphblas.core.operator.SelectOp() :members: + +Aggregator +~~~~~~~~~~ + +.. autoclass:: graphblas.core.operator.Aggregator() + :members: diff --git a/docs/api_reference/types.rst b/docs/api_reference/types.rst new file mode 100644 index 000000000..e8187caa5 --- /dev/null +++ b/docs/api_reference/types.rst @@ -0,0 +1,22 @@ +Types +----- + +DataType +~~~~~~~~ + +The object returned by :func:`~graphblas.dtypes.register_new` and +:func:`~graphblas.dtypes.register_anonymous`, and carried by every collection's +``.dtype``. The JIT introspection properties below report what SuiteSparse +actually registered for a user-defined type; see :doc:`../user_guide/udt` for +the operator-side counterparts on a typed operator (``op.jit_c_name``, +``op.jit_c_source``). + +.. autoclass:: graphblas.dtypes.DataType() + :members: jit_c_name, jit_c_definition + +Registering user-defined types +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. autofunction:: graphblas.dtypes.register_new + +.. autofunction:: graphblas.dtypes.register_anonymous diff --git a/docs/api_reference/utilities.rst b/docs/api_reference/utilities.rst new file mode 100644 index 000000000..8e6e10486 --- /dev/null +++ b/docs/api_reference/utilities.rst @@ -0,0 +1,13 @@ +Utilities +--------- + +Initialization +~~~~~~~~~~~~~~~ + +.. autofunction:: graphblas.init + +Recorder +~~~~~~~~ + +.. autoclass:: graphblas.Recorder + :members: diff --git a/docs/user_guide/operators.rst b/docs/user_guide/operators.rst index 6a66c295c..bd56ca889 100644 --- a/docs/user_guide/operators.rst +++ b/docs/user_guide/operators.rst @@ -149,6 +149,8 @@ Common semirings are: - **min_second** - **max_first** - **max_second** + - **plus_first** (sum the left operand's values over the connection pattern) + - **plus_second** (sum the right operand's values over the connection pattern) - **plus_min** - **lor_land** - **land_lor** @@ -156,6 +158,42 @@ Common semirings are: Semirings are located in the ``graphblas.semiring`` namespace. Additional semirings registered from numpy are located in ``graphblas.semiring.numpy``. +The ``first`` and ``second`` semirings +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The ``first`` binary operator returns its left input and ``second`` returns its right input, +each ignoring the other. Combined with the ``plus`` monoid they give two semirings that read +the values of only one input during a multiply: + + - ``plus_first``: for ``C << A.mxm(B, semiring.plus_first)``, each ``C[i, k]`` is the sum of + ``A[i, j]`` over the ``j`` where both ``A[i, j]`` and ``B[j, k]`` are present. ``B`` + contributes only its structure; its stored values are never read. + - ``plus_second``: the mirror image. Each ``C[i, k]`` sums ``B[j, k]`` over those same ``j``, + and ``A`` contributes only its structure. + +These are the natural choice when one operand is a boolean or *iso* (single-valued) adjacency +matrix, where a ``plus_times`` multiply either multiplies by one (a no-op) or rescales every +term by the same constant. ``first`` and ``second`` skip the product and accumulate one side's +values directly over the connection pattern. On a 0/1 matrix, ``plus_first`` counts the +length-two paths between each pair of nodes. + +.. code-block:: python + + from graphblas import Matrix, semiring + + # A: 0->1 (10), 0->2 (20); B: 1->0 (3), 2->0 (7) + # A @ B reaches node 0 from node 0 through j = 1 and j = 2. + A = Matrix.from_coo([0, 0], [1, 2], [10, 20], nrows=3, ncols=3) + B = Matrix.from_coo([1, 2], [0, 0], [3, 7], nrows=3, ncols=3) + + A.mxm(B, semiring.plus_first).new() # C[0, 0] == 30 (10 + 20, taken from A) + A.mxm(B, semiring.plus_second).new() # C[0, 0] == 10 ( 3 + 7, taken from B) + A.mxm(B, semiring.plus_times).new() # C[0, 0] == 170 (10*3 + 20*7) + +The same ``first`` / ``second`` choice pairs with other monoids: ``min_first`` and ``min_second`` +(listed above) carry a value along an edge without combining it, the pattern behind +label-propagation traversals. + IndexUnary Operators -------------------- diff --git a/docs/user_guide/udf.rst b/docs/user_guide/udf.rst index 670cbc6e3..2f5bdcc01 100644 --- a/docs/user_guide/udf.rst +++ b/docs/user_guide/udf.rst @@ -188,3 +188,27 @@ time: # numba.njit has run. The compile happens on first lookup (``unary.heavy_op[int]``). + +Compilation caching +------------------- + +A handful of built-in operators are themselves UDFs (``binary.floordiv``, +``rfloordiv``, ``absfirst``, ``abssecond``, and ``rpow``). Because these are +plain module-level functions, python-graphblas compiles them with Numba's +``cache=True``, so the Numba half of the compile is paid once per machine +instead of once per Python process. Numba writes the cache (``.nbi`` index and +``.nbc`` data files) into the ``__pycache__`` directory next to the operator +module. A warm cache shaves roughly 20% off a single op's first-touch build and +about 45% when several are used; the rest is Numba/LLVM startup and per-process +object linking, which caching cannot remove. + +If that ``__pycache__`` directory is not writable (a read-only install, for +example), Numba silently falls back to its per-user cache directory +(``~/.cache/numba`` on Linux, ``~/Library/Caches/numba`` on macOS; set +``NUMBA_CACHE_DIR`` to override). No warning is emitted for the fallback. + +Operators you create with ``register_new`` or ``register_anonymous`` are never +passed ``cache=True``: Numba can only cache a function with a stable on-disk +source, which a lambda or an interactively defined function does not have. A +registered operator's compile is paid once per process, and ``lazy=True`` is +the tool for deferring it. diff --git a/docs/user_guide/udt.rst b/docs/user_guide/udt.rst index 18af2cf3e..1d179f243 100644 --- a/docs/user_guide/udt.rst +++ b/docs/user_guide/udt.rst @@ -173,10 +173,74 @@ For nested record UDTs the tuple is *flat over the leaves*. Given return ``(id, x, y)``, not ``(id, (x, y))``. Returning an existing record value (e.g., one of the inputs) is also fine and preserves the nested shape. +For array UDTs each operand arrives as a numpy view of that element's values, in +the UDT's declared shape: a ``np.dtype((np.float64, (2, 4)))`` UDT hands the UDF +a 2-by-4 array, indexable as ``x[i, j]``. Array expressions work as written, and +the UDF may return one of its operands or build a new array that fills the +element, either at the element's own shape or at one that broadcasts to it (a +``(1,)`` return fills every slot of a ``(6,)`` element): + +.. code-block:: python + + def midpoint(x, y): + return (x + y) / 2 + + op = binary.register_new("midpoint", midpoint, is_udt=True) + + a = Vector(point3, size=1) + a[0] = [0.0, 2.0, 4.0] + b = Vector(point3, size=1) + b[0] = [10.0, 20.0, 30.0] + + c = a.ewise_mult(b, op[point3]).new() + # c[0] = [5.0, 11.0, 17.0] + +A return that cannot fill the element, such as ``x[:2]`` from a 3-element UDT, +is rejected when the op is typed for the UDT. The shape is learned by running +the UDF once on sample values, so a UDF that raises on those values is not +checked. + If your UDF references a field that doesn't exist, or returns the wrong arity, you'll get a ``UdfParseError`` with the actionable diagnostic line surfaced from Numba's typing pass instead of a 200-line traceback. +Naming the output type +~~~~~~~~~~~~~~~~~~~~~~ + +By default the return type is worked out from what the UDF returns, matched +against the input dtypes. That can only name a type the operator already has +in hand, so an output UDT that is not one of the operands is out of reach. +The ``x[:2]`` rejection above is a case of this: the output really is a +2-element UDT, but nothing says so. + +Pass ``ret_dtype`` to say it outright. It takes anything ``lookup_dtype`` +accepts and requires ``is_udt=True``:: + + nine = gb.dtypes.register_anonymous(np.dtype((np.float64, (9,))), "Nine") + three = gb.dtypes.register_anonymous(np.dtype((np.float64, (3,))), "Three") + + head = gb.core.operator.UnaryOp.register_anonymous( + lambda x: x[:3], "head", is_udt=True, ret_dtype=three + ) + head[nine].return_type # Three + +``ret_dtype`` is available on ``register_anonymous`` and ``register_new`` for +``UnaryOp``, ``BinaryOp``, ``IndexUnaryOp``, and ``IndexBinaryOp``. It is a +property of the operator, not of a particular input dtype: the same output +type applies to every dtype the operator is typed for. An operator whose +output type should vary with its inputs still needs one registration per +output type. + +Declaring the type does not switch off the shape check described above; it +points the check at the declared type instead. A UDF whose result cannot +fill a ``ret_dtype`` element is still rejected when the op is typed, so a +wrong ``ret_dtype`` is a registration error rather than a silently +mis-typed result. + +It does not apply to ``SelectOp``, whose return type GraphBLAS fixes at +``BOOL``, nor to builtin (non-UDT) dtypes, where the return type comes from +compiling the function against each input type in turn. + .. _udt_jit_introspection: JIT and introspection diff --git a/graphblas/__init__.py b/graphblas/__init__.py index 86759f570..0868a45c8 100644 --- a/graphblas/__init__.py +++ b/graphblas/__init__.py @@ -25,7 +25,45 @@ def get_config(): import donfig import yaml - config = donfig.Config("graphblas") + class Config(donfig.Config): + """donfig Config that rejects silent attribute writes. + + Options are set with ``config.set(name=value)`` (optionally as a + ``with config.set(name=value):`` block) and read with ``config[name]``. + Plain attribute assignment (``config.name = value``) would otherwise + create a dead instance attribute and leave the real option unchanged, + so we raise instead. + """ + + def __init__(self, *args, **kwargs): + # Writes are unrestricted until donfig's own __init__ finishes, and + # the attributes it set there become the allowlist. Deriving that + # list beats hardcoding donfig's internals: a donfig release that + # adds or renames a field can't then break attribute access here. + object.__setattr__(self, "_initializing", True) + super().__init__(*args, **kwargs) + object.__setattr__(self, "_donfig_attrs", frozenset(self.__dict__) - {"_initializing"}) + object.__setattr__(self, "_initializing", False) + + def __setattr__(self, key, value): + # An instance built without running __init__ has no allowlist yet, + # so stay permissive rather than raising a confusing AttributeError + # from the guard. + if self.__dict__.get("_initializing", True) or key in self._donfig_attrs: + object.__setattr__(self, key, value) + return + if key in self.config: + raise AttributeError( + f"Cannot set config option {key!r} by attribute assignment. " + f"Use `graphblas.config.set({key}=...)` to change it (optionally " + f"in a `with` block for a scoped change) and " + f"`graphblas.config[{key!r}]` to read it." + ) + raise AttributeError( + f"Unknown config option {key!r}; known options are {sorted(self.config)}." + ) + + config = Config("graphblas") path = Path(__file__).parent / "graphblas.yaml" with path.open() as f: defaults = yaml.safe_load(f) @@ -36,6 +74,10 @@ def get_config(): config = get_config() del get_config +# None until a backend is initialized. Touching a special attribute such as +# gb.Matrix auto-initializes, as does an explicit init(), and sets this to +# "suitesparse" or "suitesparse-vanilla". Reading gb.backend is not itself a +# special-attribute access, so it never triggers initialization. backend = None _init_params = None _SPECIAL_ATTRS = { diff --git a/graphblas/agg/__init__.py b/graphblas/agg/__init__.py index da7c13591..d25017279 100644 --- a/graphblas/agg/__init__.py +++ b/graphblas/agg/__init__.py @@ -109,7 +109,9 @@ def __getattr__(key): ss = import_module(".ss", __name__) globals()["ss"] = ss return ss - raise AttributeError(f"module {__name__!r} has no attribute {key!r}") + from ..core.utils import _module_attr_error + + raise _module_attr_error(__name__, key, __dir__()) from ..core import operator # noqa: E402 isort:skip diff --git a/graphblas/binary/__init__.py b/graphblas/binary/__init__.py index 1b8985f73..4c2b7a267 100644 --- a/graphblas/binary/__init__.py +++ b/graphblas/binary/__init__.py @@ -66,7 +66,9 @@ def __getattr__(key): f"module {__name__!r} unable to compile UDF for {key!r}; " "install numba for UDF support" ) - raise AttributeError(f"module {__name__!r} has no attribute {key!r}") + from ..core.utils import _module_attr_error + + raise _module_attr_error(__name__, key, __dir__()) from ..core import operator # noqa: E402 isort:skip diff --git a/graphblas/core/automethods.py b/graphblas/core/automethods.py index 600a6e139..32f6c37ca 100644 --- a/graphblas/core/automethods.py +++ b/graphblas/core/automethods.py @@ -10,11 +10,38 @@ from .. import config +# Scalar read-only attributes whose value is a plain read of the underlying C +# scalar. For an index-extract expression (v[i], A[i, j]) these resolve with a +# single extractElement straight into a cscalar via ScalarIndexExpr._extract_fast, +# skipping the extra GrB_Scalar round-trip that `.new()` + Scalar.value performs. +# Only ScalarIndexExpr defines the hook, and no Vector/Matrix expression routes +# these attrs through _get_value, so their resolution short-circuits on the +# membership test. +_fast_scalar_attrs = frozenset( + { + "value", + "__float__", + "__int__", + "__complex__", + "__index__", + "__bool__", + "__array__", + "is_empty", + "_is_empty", + } +) + def _get_value(self, attr=None, default=None): if config.get("autocompute"): if self._value is None: - self._value = self.new() + if ( + attr in _fast_scalar_attrs + and (extract_fast := getattr(self, "_extract_fast", None)) is not None + ): + self._value = extract_fast() + else: + self._value = self.new() if attr is None: return self._value return getattr(self._value, attr) @@ -339,11 +366,22 @@ def __ixor__(self, other): # End auto-generated code -def _main(): +def _main(base_dir=None, callblack=True): + # base_dir lets `scripts/autogenerate.py --check` redirect reads and writes to a + # scratch copy of the tree; when None we regenerate the real files in place. + # callblack=False keeps that check hermetic: black is optional here, and letting it + # run would make the scratch output depend on whether it is installed. + import functools from pathlib import Path from .utils import _autogenerate_code + if not callblack: + _autogenerate_code = functools.partial(_autogenerate_code, callblack=False) + + if base_dir is None: + base_dir = Path(__file__).parent + common = { "_name_html", "_nvals", @@ -461,7 +499,7 @@ def _main(): f' raise TypeError(f"{name!r} not supported for {{type(self).__name__}}")\n\n' ) - _autogenerate_code(Path(__file__), "\n".join(lines)) + _autogenerate_code(base_dir / "automethods.py", "\n".join(lines)) # Copy to scalar.py and infix.py lines = [] @@ -482,7 +520,7 @@ def _main(): continue lines.append(f" {name} = automethods.{name}") - thisdir = Path(__file__).parent + thisdir = base_dir infix_exclude = {"_get_value"} def get_name(line): diff --git a/graphblas/core/base.py b/graphblas/core/base.py index 15f66bc2f..ae9c1fbca 100644 --- a/graphblas/core/base.py +++ b/graphblas/core/base.py @@ -2,12 +2,11 @@ from .. import backend, config from .. import replace as replace_singleton -from ..dtypes import BOOL from ..exceptions import check_status from . import NULL from .descriptor import lookup as descriptor_lookup from .expr import AmbiguousAssignOrExtract, Updater -from .mask import Mask +from .mask import Mask, _check_mask from .operator import UNKNOWN_OPCLASS, binary_from_string, find_opclass, get_typed_op from .utils import _Pointer, libget, output_type @@ -20,6 +19,11 @@ def record_raw(text): rec.record_raw(text) +def _is_recording(): + """Whether a Recorder is active; fast paths that bypass ``call`` must check this.""" + return _recorder.get(_prev_recorder) is not None + + def call(cfunc_name, args): call_args = [getattr(x, "_carg", x) if x is not None else NULL for x in args] cfunc = libget(cfunc_name) @@ -166,21 +170,19 @@ def _expect_op(self, op, values, *, within, **kwargs): AmbiguousAssignOrExtract._expect_type = _expect_type -def _check_mask(mask, output=None): - if not isinstance(mask, Mask): - # Convert bool objects to value masks - if output_type(mask).__name__ in {"Vector", "Matrix"}: - if mask.dtype != BOOL: - raise TypeError( - f"Mask must be boolean objects (got {mask.dtype}) " - "or indicate values (M.V) or structure (M.S)" - ) - mask = mask.V # auto-compute (will raise if disabled) - else: - raise TypeError(f"Invalid mask: {type(mask)}") - if output is not None and output.ndim == 1 and mask.parent.ndim != 1: - raise TypeError(f"Mask object must be type Vector; got {type(mask.parent)}") - return mask +# Curated hints for common attribute-access mistakes on Vector/Matrix/Scalar. +# Only consulted from __getattr__, which fires solely on a genuine attribute +# miss, so the normal (slotted) attribute hot path is untouched. +_INSTANCE_ATTR_HINTS = { + # `.new()` resolves expressions (e.g. `A.mxm(B).new()`); a concrete + # object is copied with `.dup()`. + "new": ( + "`.new()` resolves an expression (e.g. `A.mxm(B).new()`); a concrete " + "object has no `.new()`. Use `.dup()` to copy this object." + ), + # transpose is the `.T` property, not a method. + "transpose": "transpose is the `.T` property, e.g. `A.T` (not a method call).", +} class BaseType: @@ -189,6 +191,14 @@ class BaseType: # Flag for operations which depend on scalar vs vector/matrix _is_scalar = False + def __getattr__(self, name): + # Fires only on a genuine attribute miss (slots/methods resolve first), + # so this is free on the hot path. Adds hints for common mistakes. + base = f"{type(self).__name__!r} object has no attribute {name!r}" + if (hint := _INSTANCE_ATTR_HINTS.get(name)) is not None: + raise AttributeError(f"{base}; {hint}") + raise AttributeError(base) + def __call__( self, *optional_mask_accum_replace, @@ -460,7 +470,11 @@ def _update(self, expr, mask=None, accum=None, replace=False, input_mask=None, * complement = False structure = False else: - mask = _check_mask(mask, self) + # Assignment (`method_name == "__setitem__"`) may target a Matrix + # row/column with a Vector mask, so only enforce the strict + # mask-kind match for full-tensor operations. + strict_kind = expr.method_name != "__setitem__" + mask = _check_mask(mask, self, strict_kind=strict_kind) complement = mask.complement structure = mask.structure @@ -611,7 +625,9 @@ def _new(self, dtype, mask, name, is_cscalar=None, **opts): elif mask is None: output.update(self, **opts) else: - mask = _check_mask(mask, output) + # `.new()` always builds a full output matching this expression, so + # the mask kind must match the output dimensions. + mask = _check_mask(mask, output, strict_kind=True) output(mask=mask, **opts).update(self) return output diff --git a/graphblas/core/dtypes.py b/graphblas/core/dtypes.py index 8eb09725b..086ce7abf 100644 --- a/graphblas/core/dtypes.py +++ b/graphblas/core/dtypes.py @@ -549,6 +549,27 @@ def lookup_dtype(key, value=None): raise ValueError(f"Unknown dtype: {key} of type {type(key)}") +def _raise_dtype_or_arraylike(cls_name, dtype, exc): + """Turn a failed constructor dtype lookup into a helpful error. + + Call from ``Vector``/``Matrix`` after ``lookup_dtype(dtype)`` has already + raised ``exc``. The constructors take a dtype as the first argument, so a + common mistake is to pass the data instead (``Vector([1, 2, 3])``). When + ``dtype`` is array-like data rather than a dtype spec, point at the ``from_*`` + constructors; otherwise re-raise the original ``Unknown dtype`` error. Valid + list/tuple dtype specs (structured and subarray dtypes) never reach here + because ``lookup_dtype`` accepts them. + """ + if isinstance(dtype, (list, tuple, np.ndarray)): + raise TypeError( + f"{cls_name}() expects a dtype as the first argument, not " + f"{type(dtype).__name__} data. To build a {cls_name} from existing " + f"values, use a constructor such as {cls_name}.from_coo(...) or " + f"{cls_name}.from_dense(...)." + ) from None + raise exc + + def unify(type1, type2, *, is_left_scalar=False, is_right_scalar=False): """Returns a type that can hold both type1 and type2. diff --git a/graphblas/core/expr.py b/graphblas/core/expr.py index 47ff18a0b..e3220758f 100644 --- a/graphblas/core/expr.py +++ b/graphblas/core/expr.py @@ -176,6 +176,21 @@ def parse_indices(self, indices, shape): def parse_index(self, index, typ, size): from .scalar import _as_scalar + if typ is int: + # Fast lane for a plain Python int, the overwhelmingly common case. + # The two np.issubdtype checks below cost ~125ns each (measured), a + # meaningful slice of single-element index parsing. output_type maps + # only the exact `int` type to `int` (bool -> bool, numpy ints -> their + # own type), so bool and numpy scalars never enter here and keep their + # existing handling. A plain int is always signed, so the negative + # branch always applies, matching the signedinteger branch below. + if index >= size: + raise IndexError(f"Index out of range: index={index}, size={size}") + if index < 0: + index = index + size + if index < 0: + raise IndexError(f"Index out of range: index={index - size}, size={size}") + return AxisIndex(None, _as_scalar(index, _INDEX, is_cscalar=True), None, size) if np.issubdtype(typ, np.integer): if index >= size: raise IndexError(f"Index out of range: index={index}, size={size}") @@ -321,7 +336,7 @@ def new(self, dtype=None, *, mask=None, input_mask=None, name=None, **opts): if input_mask is not None: if mask is not None: raise TypeError("mask and input_mask arguments cannot both be given") - from .base import _check_mask + from .mask import _check_mask input_mask = _check_mask(input_mask, self.parent) mask = self._input_mask_to_mask(input_mask, **opts) diff --git a/graphblas/core/formatting.py b/graphblas/core/formatting.py index 5fe9b6972..cd9380c0c 100644 --- a/graphblas/core/formatting.py +++ b/graphblas/core/formatting.py @@ -1,4 +1,11 @@ -# This file imports pandas, so it should only be imported when formatting +# The rich repr and _repr_html_ are hand-rendered here to reproduce pandas' +# DataFrame text/HTML output byte-for-byte without importing pandas. When pandas +# is installed we still read its display.* options so a user's option_context is +# honored; when it is absent we fall back to pandas' documented defaults. +import math +import re +import shutil + import numpy as np from .. import backend, config, monoid, unary @@ -14,6 +21,35 @@ except ImportError: # pragma: no cover (import) has_pandas = False +# pandas display.* defaults, used verbatim when pandas is not installed so the +# hand renderer produces the same output it would with a freshly imported pandas. +_DISPLAY_DEFAULTS = { + "max_rows": 60, + "min_rows": 10, + "max_columns": 0, + "width": 80, + "expand_frame_repr": True, + "precision": 6, + "max_colwidth": 50, + "chop_threshold": None, + "colheader_justify": "right", + "float_format": None, + "html.border": 1, + "html.use_mathjax": True, +} + + +def _display_option(name): + """Read a pandas display. option, or fall back to its default. + + Reading from pandas keeps a user's ``pd.option_context`` honored while the + rendering logic itself stays pandas-free. + """ + if has_pandas: + return pd.get_option(f"display.{name}") + return _DISPLAY_DEFAULTS[name] + + # This was written by a complete novice at CSS. # If you can help make it better, please do! CSS_STYLE = """ @@ -219,7 +255,7 @@ def _update_vector_array(arr, vector, columns, column_offset, *, mask=None): def _get_max_columns(): - max_columns = pd.options.display.max_columns + max_columns = _display_option("max_columns") if max_columns == 0: # We are probably in a terminal and pandas will automatically size the data correctly. # In this case, let's get a sufficiently large amount of data to show and defer to pandas. @@ -241,13 +277,647 @@ def _get_chunk(length, min_length, max_length): return chunk, chunk_groups +class _Column: + """One rendered column: a label, its raw cell values, and how to format them. + + ``kind`` selects the pandas array formatter this column would have used: + "object" (GenericArrayFormatter), "float" (FloatArrayFormatter, also complex), + or "int" (IntArrayFormatter). ``numeric`` mirrors pandas ``is_numeric_dtype`` + and controls whether the column header gets a leading space. + """ + + __slots__ = ("label", "values", "kind", "numeric") + + def __init__(self, label, values, kind, numeric): + self.label = label + self.values = values + self.kind = kind + self.numeric = numeric + + +class _GBFrame: + """A minimal object/typed-column table, standing in for a pandas DataFrame. + + ``col_name`` is the columns' index name (pandas ``df.columns.name``); when set + it appears in the top-left corner cell, as vector reprs rely on. + """ + + __slots__ = ("columns", "index", "col_name") + + def __init__(self, columns, index, col_name=None): + self.columns = columns + self.index = index + self.col_name = col_name + + @property + def ncols(self): + return len(self.columns) + + @property + def nrows(self): + return len(self.index) + + def slice_cols(self, idx): + return _GBFrame([self.columns[i] for i in idx], self.index, self.col_name) + + def slice_rows(self, idx): + columns = [ + _Column(c.label, [c.values[i] for i in idx], c.kind, c.numeric) for c in self.columns + ] + return _GBFrame(columns, [self.index[i] for i in idx], self.col_name) + + +def _isna_cell(x): + # Gaps in the dense grid are float NaN; present values (including displayed + # "nan"/"inf") are never a bare float NaN, so this only flags the gaps. + return x is None or (isinstance(x, float) and math.isnan(x)) + + +def _count_present(arr): + return sum(1 for x in arr.flat if not _isna_cell(x)) + + +def _dtype_kind_numeric(dtype): + kind = dtype.kind + if kind in "fc": + return "float", True + if kind in "iu": + return "int", True + if kind == "b": + # bool renders via the generic formatter, but is_numeric_dtype(bool) is True + return "object", True + return "object", False + + +def _make_dense_frame(arr, columns, index): + nrows = len(index) + out = [] + for j, label in enumerate(columns): + vals = [("" if _isna_cell(arr[i, j]) else arr[i, j]) for i in range(nrows)] + out.append(_Column(label, vals, "object", False)) + return _GBFrame(out, list(index)) + + +def _make_coo_frame(label_arrays, add_dots): + n = len(label_arrays[0][1]) + index = list(range(n)) + out = [] + for label, values in label_arrays: + values = np.asarray(values) + if add_dots: + out.append(_Column(label, [*values.tolist(), "..."], "object", False)) + else: + kind, numeric = _dtype_kind_numeric(values.dtype) + out.append(_Column(label, values.tolist(), kind, numeric)) + if add_dots: + index.append("...") + return _GBFrame(out, index) + + +# --- cell formatting (reproduces pandas array formatters for object/int/float) --- + +_NUMBER_RE = re.compile(r"^\s*[\+-]?[0-9]+\.[0-9]*$") + + +def _is_float_scalar(v): + # Matches pandas.lib.is_float: python/numpy floats, but not bool/int/complex. + return isinstance(v, (float, np.floating)) + + +def _pprint(v): + # Reproduces pandas printing.pprint_thing for our cell types (escape_chars for + # tab/cr/nl, quote_strings=False): scalars -> str, sequences recurse. + if isinstance(v, (list, tuple)): + body = ", ".join(_pprint(e) for e in v) + if isinstance(v, tuple) and len(v) == 1: + body += "," + return f"[{body}]" if isinstance(v, list) else f"({body})" + s = str(v) + return s.replace("\t", r"\t").replace("\r", r"\r").replace("\n", r"\n") + + +def _trim_zeros_single_float(s): + s = s.rstrip("0") + if s.endswith("."): + s += "0" + return s + + +def _trim_zeros_float(str_floats): + trimmed = list(str_floats) + + def is_number_with_decimal(x): + return _NUMBER_RE.match(x) is not None + + def should_trim(values): + numbers = [x for x in values if is_number_with_decimal(x)] + return len(numbers) > 0 and all(x.endswith("0") for x in numbers) + + while should_trim(trimmed): + trimmed = [x[:-1] if is_number_with_decimal(x) else x for x in trimmed] + return [x + "0" if is_number_with_decimal(x) and x.endswith(".") else x for x in trimmed] + + +def _trim_zeros_complex(str_complexes): + real_part, imag_part = [], [] + for x in str_complexes: + trimmed = re.split(r"(?{padded_length}}" + "j" + for real_pt, imag_pt in zip(padded_parts[:n], padded_parts[n:], strict=True) + ] + + +def _value_formatter(fmt_str, threshold): + def base(v): + return fmt_str.format(value=v) + + if threshold is None: + return base + + def formatter(v): + return base(v) if abs(v) > threshold else base(0.0) + + return formatter + + +def _format_reals_with_na(values, formatter, na_rep): + return [na_rep if (v != v) else formatter(v) for v in values] + + +def _format_complex_with_na(values, formatter, na_rep): + out = [] + for val in values: + re_v, im_v = val.real, val.imag + re_na, im_na = re_v != re_v, im_v != im_v + if not re_na and not im_na: + out.append(formatter(val)) + elif not re_na: + out.append(f"{formatter(re_v)}+{na_rep}j") + elif not im_na: + imag_formatted = formatter(im_v).strip() + if imag_formatted.startswith("-"): + out.append(f"{na_rep}{imag_formatted}j") + else: + out.append(f"{na_rep}+{imag_formatted}j") + else: + out.append(f"{na_rep}+{na_rep}j") + return out + + +def _format_float_column(values, digits): + # Reproduces FloatArrayFormatter for fixed_width, leading_space=True, na_rep="NaN". + arr = np.asarray(values) + is_complex = np.iscomplexobj(arr) + na_rep = "NaN" + if (float_format := _display_option("float_format")) is not None: + # A user display.float_format callable makes FloatArrayFormatter drop + # fixed_width: each value (real or complex) is just float_format(value), + # with no trailing-zero trim and no scientific switchover. Iterate the + # numpy array (not .tolist()) so the callable receives numpy scalars, as + # pandas does; e.g. a "%f"-style callable then casts complex the same way. + return [na_rep if (v != v) else float_format(v) for v in arr] + seq = arr.tolist() + threshold = _display_option("chop_threshold") + + def format_with(fmt_str): + formatter = _value_formatter(fmt_str, threshold) + if is_complex: + return _trim_zeros_complex(_format_complex_with_na(seq, formatter, na_rep)) + return _trim_zeros_float(_format_reals_with_na(seq, formatter, na_rep)) + + result = format_with(f"{{value: .{digits:d}f}}") + too_long = bool(result) and max(len(x) for x in result) > digits + 6 + abs_vals = np.abs(arr) + has_large = bool((abs_vals > 1e6).any()) + has_small = bool(((abs_vals < 10.0 ** (-digits)) & (abs_vals > 0)).any()) + if has_small or (too_long and has_large): + result = format_with(f"{{value: .{digits:d}e}}") + return list(result) + + +def _justify(strings, width, mode="right"): + if mode == "left": + return [x.ljust(width) for x in strings] + if mode == "center": + return [x.center(width) for x in strings] + return [x.rjust(width) for x in strings] + + +def _make_fixed_width(strings, justify="right", minimum=None): + if not strings: + return list(strings) + max_len = max(len(x) for x in strings) + if minimum is not None: + max_len = max(minimum, max_len) + conf_max = _display_option("max_colwidth") + if conf_max is not None and max_len > conf_max: + max_len = conf_max + + def just(x): + if conf_max is not None and conf_max > 3 and len(x) > max_len: + x = x[: max_len - 3] + "..." + return x + + return _justify([just(x) for x in strings], max_len, justify) + + +def _format_labels(labels): + # Reproduce pandas Index._format_flat(include_name=False) for our label types: + # integer labels are padded to a uniform width (left-justified, with a sign + # column when any are negative); string labels are left as-is. + if labels and all(isinstance(x, (int, np.integer)) and not isinstance(x, bool) for x in labels): + pattern = "{: d}" if any(x < 0 for x in labels) else "{:d}" + strs = [pattern.format(x) for x in labels] + width = max(len(s) for s in strs) + return [s.ljust(width) for s in strs] + return [str(x) for x in labels] + + +def _adjoin(space, lists): + # Port of pandas printing.adjoin (ascii len/ljust); glues columns with `space`. + lengths = [max(map(len, x)) + space for x in lists[:-1]] + lengths.append(max(map(len, lists[-1]))) + max_len = max(map(len, lists)) + padded = [] + for i, lst in enumerate(lists): + nl = [x.ljust(lengths[i]) for x in lst] + nl = [" " * lengths[i]] * (max_len - len(lst)) + nl + padded.append(nl) + return "\n".join("".join(parts) for parts in zip(*padded, strict=True)) + + +def _binify(cols, line_width): + adjoin_width = 1 + bins = [] + curr_width = 0 + i_last = len(cols) - 1 + for i, w in enumerate(cols): + w_adjoined = w + adjoin_width + curr_width += w_adjoined + if i_last == i: + wrap = curr_width + 1 > line_width and i > 0 + else: + wrap = curr_width + 2 > line_width and i > 0 + if wrap: + bins.append(i) + curr_width = w_adjoined + bins.append(len(cols)) + return bins + + +def _console_width(): + # pandas repr sets the wrap width from console.get_console_size(); reuse it + # when present so the wrap decision is identical. Without pandas (or if that + # private module moved) fall back to display.width. + if has_pandas: + try: + from pandas.io.formats.console import get_console_size + + return get_console_size()[0] + except Exception: # pragma: no cover (defensive across pandas versions) + pass + return _display_option("width") + + +class _TextFormatter: + """Reproduces pandas DataFrameFormatter + StringFormatter for text repr.""" + + def __init__(self, frame, max_rows, min_rows, max_cols): + self.frame = frame + self.max_rows = max_rows + self.min_rows = min_rows + self.max_cols = max_cols + self.justify = _display_option("colheader_justify") + self.tr_frame = frame + self.tr_col_num = None + self.tr_row_num = None + self.max_cols_fitted = self._calc_max_cols_fitted() + self.max_rows_fitted = self._calc_max_rows_fitted() + self.truncate() + + def _is_in_terminal(self): + return self.max_cols == 0 or self.max_rows == 0 + + def _calc_max_cols_fitted(self): + if not self._is_in_terminal(): + return self.max_cols + width = shutil.get_terminal_size()[0] + if self.max_cols == 0 and self.frame.ncols > width: + return width + return self.max_cols + + def _calc_max_rows_fitted(self): + if self._is_in_terminal() and self.max_rows == 0: + # rows available for data: terminal height minus dots + prompt + header + return shutil.get_terminal_size()[1] - 3 + max_rows = self.max_rows + if max_rows and self.frame.nrows > max_rows and self.min_rows: + max_rows = min(self.min_rows, max_rows) + return max_rows + + @property + def is_truncated_horizontally(self): + return bool(self.max_cols_fitted and self.frame.ncols > self.max_cols_fitted) + + @property + def is_truncated_vertically(self): + return bool(self.max_rows_fitted and self.frame.nrows > self.max_rows_fitted) + + @property + def is_truncated(self): + return self.is_truncated_horizontally or self.is_truncated_vertically + + def truncate(self): + if self.is_truncated_horizontally: + self._truncate_horizontally() + if self.is_truncated_vertically: + self._truncate_vertically() + + def _truncate_horizontally(self): + col_num = self.max_cols_fitted // 2 + if col_num >= 1: + _len = self.tr_frame.ncols + self.tr_frame = self.tr_frame.slice_cols( + [*range(col_num), *range(_len - col_num, _len)] + ) + else: + col_num = self.max_cols + self.tr_frame = self.tr_frame.slice_cols(list(range(col_num))) + self.tr_col_num = col_num + + def _truncate_vertically(self): + row_num = self.max_rows_fitted // 2 + if row_num >= 1: + _len = self.tr_frame.nrows + self.tr_frame = self.tr_frame.slice_rows( + [*range(row_num), *range(_len - row_num, _len)] + ) + else: + row_num = self.max_rows + self.tr_frame = self.tr_frame.slice_rows(list(range(row_num))) + self.tr_row_num = row_num + + def _format_col_raw(self, col): + if col.kind == "int": + return [f"{x: d}" for x in col.values] + if col.kind == "float": + return _format_float_column(col.values, _display_option("precision")) + precision = _display_option("precision") + float_format = _display_option("float_format") + out = [] + for v in col.values: + # A float NaN is excluded from pandas' float-format branch (it uses + # is_float(v) & notna(v)) and rendered as the na_rep "NaN" instead. + if _is_float_scalar(v) and not math.isnan(v): + if float_format is not None: + # A user display.float_format callable replaces the default + # precision render (and adds no sign-space of its own). + out.append(float_format(v)) + else: + out.append(_trim_zeros_single_float(f"{v: .{precision}f}")) + elif v is None: + out.append(" None") + elif _is_float_scalar(v): + out.append(" NaN") + else: + out.append(f" {_pprint(v)}") + return out + + def _get_body_strcols(self): + # Column labels are formatted together (integer labels padded to a uniform + # width) the way pandas Index._format_flat does, not per column. + labels = _format_labels([col.label for col in self.tr_frame.columns]) + strcols = [] + for col, label in zip(self.tr_frame.columns, labels, strict=True): + header = f" {label}" if col.numeric else label + header_colwidth = len(header) + # pandas fixes width twice: format_array right-justifies to the cell + # content width, then the body pass re-justifies with colheader_justify + # (which only matters when the header is wider, or when it is "left"). + fmt_values = _make_fixed_width(self._format_col_raw(col), "right") + fmt_values = _make_fixed_width(fmt_values, self.justify, minimum=header_colwidth) + max_len = max(max((len(x) for x in fmt_values), default=0), header_colwidth) + cheader = _justify([header], max_len, self.justify) + strcols.append(cheader + fmt_values) + return strcols + + def _get_index_strcol(self): + idx = _make_fixed_width([str(x) for x in self.tr_frame.index], justify="left") + corner = "" if self.frame.col_name is None else str(self.frame.col_name) + return [corner, *idx] + + def get_strcols(self): + strcols = self._get_body_strcols() + strcols.insert(0, self._get_index_strcol()) + return strcols + + @property + def _adjusted_tr_col_num(self): + return self.tr_col_num + 1 # index column is always shown + + def _insert_dot_separators(self, strcols): + index_length = len(self._get_index_strcol()) + if self.is_truncated_horizontally: + strcols.insert(self._adjusted_tr_col_num, [" ..."] * index_length) + if self.is_truncated_vertically: + self._insert_dots_vertical(strcols, index_length) + return strcols + + def _insert_dots_vertical(self, strcols, index_length): + n_header_rows = index_length - self.tr_frame.nrows + row_num = self.tr_row_num + for ix, col in enumerate(strcols): + cwidth = len(col[row_num]) + is_dot_col = self.is_truncated_horizontally and ix == self._adjusted_tr_col_num + dots = "..." if (cwidth > 3 or is_dot_col) else ".." + if ix == 0: + dot_mode = "left" + elif is_dot_col: + cwidth = 4 + dot_mode = "right" + else: + dot_mode = "right" + col.insert(row_num + n_header_rows, _justify([dots], cwidth, dot_mode)[0]) + + def _get_strcols(self): + strcols = self.get_strcols() + if self.is_truncated: + strcols = self._insert_dot_separators(strcols) + return strcols + + def _fit_to_terminal(self, strcols): + lines = _adjoin(1, strcols).split("\n") + max_len = max(len(x) for x in lines) + width = shutil.get_terminal_size()[0] + adj_dif = max_len - width + 1 # +1 to avoid too-wide repr (pandas GH #17023) + col_lens = [max((len(x) for x in col), default=0) for col in strcols] + n_cols = len(col_lens) + while adj_dif > 0 and n_cols > 1: + mid = round(n_cols / 2) + adj_dif -= col_lens.pop(mid) + 1 + n_cols = len(col_lens) + max_cols_fitted = max(n_cols - 1, 2) # minus index column; show at least two + self.max_cols_fitted = max_cols_fitted + self.truncate() + return _adjoin(1, self._get_strcols()) + + def _join_multiline(self, strcols, line_width): + adjoin_width = 1 + strcols = list(strcols) + idx = strcols.pop(0) + line_width -= max(len(x) for x in idx) + adjoin_width + col_widths = [max((len(x) for x in col), default=0) for col in strcols] + col_bins = _binify(col_widths, line_width) + nbins = len(col_bins) + blocks = [] + start = 0 + for i, end in enumerate(col_bins): + row = strcols[start:end] + row.insert(0, idx) + if nbins > 1: + nrows = len(row[-1]) + if end <= len(strcols) and i < nbins - 1: + row.append([" \\", *[" "] * (nrows - 1)]) + else: + row.append([" "] * nrows) + blocks.append(_adjoin(adjoin_width, row)) + start = end + return "\n\n".join(blocks) + + def to_string(self, line_width): + strcols = self._get_strcols() + if line_width is None: + return _adjoin(1, strcols) + if self.max_cols > 0: + return self._join_multiline(strcols, line_width) + return self._fit_to_terminal(strcols) + + +def _render_text(frame): + max_cols = _display_option("max_columns") + line_width = _console_width() if _display_option("expand_frame_repr") else None + fmt = _TextFormatter(frame, _display_option("max_rows"), _display_option("min_rows"), max_cols) + return fmt.to_string(line_width) + + +# The scoped style block pandas' NotebookFormatter emits ahead of the table. +_HTML_STYLE = ( + "" +) + + +def _html_escape(s): + return s.replace("&", "&").replace("<", "<").replace(">", ">") + + +class _HtmlBuilder: + """Reproduces pandas NotebookFormatter (DataFrame._repr_html_) markup.""" + + indent_delta = 2 + + def __init__(self, fmt): + self.fmt = fmt + self.tr = fmt.tr_frame + self.ncols = self.tr.ncols + self.th = fmt.is_truncated_horizontally + self.tv = fmt.is_truncated_vertically + self.row_levels = 1 # single-level index, always shown + self.elements = [] + + def write(self, s, indent=0): + self.elements.append(" " * indent + s) + + def _cell(self, s, kind, indent): + rs = _html_escape(str(s)).strip().replace(" ", "  ") + self.write(f"<{kind}>{rs}", indent) + + def write_tr(self, line, indent, header=False, align=None, nindex_levels=0): + self.write("" if align is None else f'', indent) + inner = indent + self.indent_delta + for i, s in enumerate(line): + self._cell(s, "th" if (header or i < nindex_levels) else "td", inner) + self.write("", indent) + + def _col_header(self, indent): + row = ["" if self.tr.col_name is None else str(self.tr.col_name)] + row.extend(_format_labels([col.label for col in self.tr.columns])) + if self.th: + row.insert(self.row_levels + self.fmt.tr_col_num, "...") + self.write_tr(row, indent, header=True, align=self.fmt.justify) + + def _body(self, indent): + index_labels = _format_labels(list(self.tr.index)) + col_cells = [ + _make_fixed_width(self.fmt._format_col_raw(col), "right") for col in self.tr.columns + ] + row = [] + for i in range(self.tr.nrows): + if self.tv and i == self.fmt.tr_row_num: + self.write_tr(["..."] * len(row), indent, nindex_levels=self.row_levels) + row = [index_labels[i], *(col_cells[j][i] for j in range(self.ncols))] + if self.th: + row.insert(self.fmt.tr_col_num + self.row_levels, "...") + self.write_tr(row, indent, nindex_levels=self.row_levels) + + def _table(self, indent=0): + classes = "dataframe" + if not _display_option("html.use_mathjax"): + classes = "dataframe tex2jax_ignore mathjax_ignore" + # pandas keeps the attribute for any non-None border, including 0. + border = _display_option("html.border") + border_attr = "" if border is None else f' border="{border}"' + self.write(f'', indent) + self.write("", indent + self.indent_delta) + self._col_header(indent + 2 * self.indent_delta) + self.write("", indent + self.indent_delta) + self.write("", indent + self.indent_delta) + self._body(indent + 2 * self.indent_delta) + self.write("", indent + self.indent_delta) + self.write("", indent) + + def render(self): + self.write("
") + self.write(_HTML_STYLE) + self._table(0) + self.write("
") + return "\n".join(self.elements) + + +def _render_html(frame): + fmt = _TextFormatter( + frame, + _display_option("max_rows"), + _display_option("min_rows"), + _display_option("max_columns"), + ) + return _HtmlBuilder(fmt).render() + + def _get_matrix_dataframe(matrix, max_rows, min_rows, max_columns, *, mask=None): - if not has_pandas: - return if max_rows is None: # pragma: no branch - max_rows = pd.options.display.max_rows + max_rows = _display_option("max_rows") if min_rows is None: # pragma: no branch - min_rows = pd.options.display.min_rows + min_rows = _display_option("min_rows") if max_columns is None: # pragma: no branch max_columns = _get_max_columns() rows, row_groups = _get_chunk(matrix._nrows, min_rows, max_rows) @@ -264,12 +934,12 @@ def _get_matrix_dataframe(matrix, max_rows, min_rows, max_columns, *, mask=None) column_offset, mask=mask, ) - df = pd.DataFrame(arr, columns=columns, index=rows) + present = _count_present(arr) + truncated = (len(rows), len(columns)) != matrix.shape if ( (mask is None or mask.structure) - and df.shape != matrix.shape - and min(matrix._nvals, max_rows if matrix._nvals <= max_rows else min_rows) - > 2 * df.count().sum() + and truncated + and min(matrix._nvals, max_rows if matrix._nvals <= max_rows else min_rows) > 2 * present ): # The data is sparse and it's better to show in COO format. # SS, SuiteSparse-specific: head @@ -283,47 +953,43 @@ def _get_matrix_dataframe(matrix, max_rows, min_rows, max_columns, *, mask=None) vals = np.zeros(vals.size, dtype=np.uint8) else: vals = np.ones(vals.size, dtype=np.uint8) - df = pd.DataFrame({"row": rows, "col": cols, "val": vals}) - if num_rows < matrix._nvals: - df.loc["..."] = ["..."] * 3 - return df - if mask is not None and not mask.structure and df.shape != matrix.shape: + return _make_coo_frame( + [("row", rows), ("col", cols), ("val", vals)], num_rows < matrix._nvals + ) + if mask is not None and not mask.structure and truncated: # This performs more calculation and uses more memory than I would prefer. # Perhaps we could use the efficient "constant vector or matrix" trick. nonzero = matrix.apply(unary.one["UINT8"]).new(mask=matrix.V, name="") num_rows = matrix._nvals if matrix._nvals <= max_rows else min_rows - if min(nonzero._nvals, num_rows) > 2 * df.count().sum(): + if min(nonzero._nvals, num_rows) > 2 * present: rows, cols, vals = nonzero.ss.head(num_rows, sort=True) if mask.complement: if not vals.flags.writeable: # pragma: no cover (safety) vals = vals.copy() vals[:] = 0 - df = pd.DataFrame({"row": rows, "col": cols, "val": vals}) - if num_rows < nonzero._nvals: - df.loc["..."] = ["..."] * 3 - return df - return df.where(pd.notna(df), "") + return _make_coo_frame( + [("row", rows), ("col", cols), ("val", vals)], num_rows < nonzero._nvals + ) + return _make_dense_frame(arr, columns, rows) def _get_vector_dataframe(vector, max_rows, min_rows, max_columns, *, mask=None): - if not has_pandas: - return if max_rows is None: # pragma: no branch - max_rows = pd.options.display.max_rows + max_rows = _display_option("max_rows") if min_rows is None: # pragma: no branch - min_rows = pd.options.display.min_rows + min_rows = _display_option("min_rows") if max_columns is None: # pragma: no branch max_columns = _get_max_columns() columns, column_groups = _get_chunk(vector._size, max_columns, max_columns) arr = np.full((1, len(columns)), np.nan, dtype=object) for column_group, column_offset in column_groups: _update_vector_array(arr, vector, column_group, column_offset, mask=mask) - df = pd.DataFrame(arr, columns=columns, index=[""]) + present = _count_present(arr) + truncated = len(columns) != vector._size if ( (mask is None or mask.structure) - and df.size != vector._size - and min(vector._nvals, max_rows if vector._nvals <= max_rows else min_rows) - > 2 * df.count().sum() + and truncated + and min(vector._nvals, max_rows if vector._nvals <= max_rows else min_rows) > 2 * present ): # The data is sparse and it's better to show in COO format. # SS, SuiteSparse-specific: head @@ -334,26 +1000,20 @@ def _get_vector_dataframe(vector, max_rows, min_rows, max_columns, *, mask=None) vals = np.zeros(vals.size, dtype=np.uint8) else: vals = np.ones(vals.size, dtype=np.uint8) - df = pd.DataFrame({"index": indices, "val": vals}) - if num_rows < vector._nvals: - df.loc["..."] = ["..."] * 2 - return df - if mask is not None and not mask.structure and df.size != vector._size: + return _make_coo_frame([("index", indices), ("val", vals)], num_rows < vector._nvals) + if mask is not None and not mask.structure and truncated: # This performs more calculation and uses more memory than I would prefer. # Perhaps we could use the efficient "constant vector or matrix" trick. nonzero = vector.apply(unary.one["UINT8"]).new(mask=vector.V, name="") num_rows = vector._nvals if vector._nvals <= max_rows else min_rows - if min(nonzero._nvals, num_rows) > 2 * df.count().sum(): + if min(nonzero._nvals, num_rows) > 2 * present: indices, vals = nonzero.ss.head(num_rows, sort=True) if mask.complement: if not vals.flags.writeable: # pragma: no cover (safety) vals = vals.copy() vals[:] = 0 - df = pd.DataFrame({"index": indices, "val": vals}) - if num_rows < nonzero._nvals: - df.loc["..."] = ["..."] * 2 - return df - return df.where(pd.notna(df), "") + return _make_coo_frame([("index", indices), ("val", vals)], num_rows < nonzero._nvals) + return _make_dense_frame(arr, columns, [""]) def get_format(x, is_transposed=False): @@ -436,14 +1096,9 @@ def vector_expression_header_html(matrix, expr): return create_header_html(name, keys, vals) -def _format_html(name, header, df, collapse): - if has_pandas: - state = "" if collapse else " open" - with pd.option_context("display.show_dimensions", False, "display.large_repr", "truncate"): - details = df._repr_html_() - else: - state = "" - details = "(Install pandas to see a preview of the data)" +def _format_html(name, header, frame, collapse): + state = "" if collapse else " open" + details = _render_html(frame) return ( "
" f"{CSS_STYLE}" @@ -667,17 +1322,12 @@ def format_matrix(matrix, *, max_rows=None, min_rows=None, max_columns=None, mas name, keys, vals, - lower_border=has_pandas, + lower_border=True, name=matrix.name if mask is None else mask.name, ) - if has_pandas: - df = _get_matrix_dataframe(matrix, max_rows, min_rows, max_columns, mask=mask) - if 0 not in matrix.shape: - with pd.option_context( - "display.show_dimensions", False, "display.large_repr", "truncate" - ): - df_repr = df.__repr__() - return f"{header}\n{df_repr}" + if 0 not in matrix.shape: + frame = _get_matrix_dataframe(matrix, max_rows, min_rows, max_columns, mask=mask) + return f"{header}\n{_render_text(frame)}" return header @@ -687,20 +1337,16 @@ def format_vector(vector, *, max_rows=None, min_rows=None, max_columns=None, mas name, keys, vals, - lower_border=has_pandas, + lower_border=True, name=vector.name if mask is None else mask.name, ) - if has_pandas: - df = _get_vector_dataframe(vector, max_rows, min_rows, max_columns, mask=mask) - if vector._size > 0: - if df.columns[0] != "index": - df.columns.name = "index" - df.index = ["value"] - with pd.option_context( - "display.show_dimensions", False, "display.large_repr", "truncate" - ): - df_repr = df.__repr__() - return f"{header}\n{df_repr}" + if vector._size > 0: + frame = _get_vector_dataframe(vector, max_rows, min_rows, max_columns, mask=mask) + if frame.columns[0].label != "index": + # Dense vectors label the corner "index" and the single row "value". + frame.col_name = "index" + frame.index = ["value"] + return f"{header}\n{_render_text(frame)}" return header diff --git a/graphblas/core/infixmethods.py b/graphblas/core/infixmethods.py index 67248223a..3f7f70ccb 100644 --- a/graphblas/core/infixmethods.py +++ b/graphblas/core/infixmethods.py @@ -320,7 +320,11 @@ def __itruediv__(self, other): # End auto-generated code -def _main(): +def _main(base_dir=None, callblack=True): + # base_dir lets `scripts/autogenerate.py --check` redirect the write to a scratch + # copy of the tree; when None we regenerate infixmethods.py in place. + # callblack=False keeps that check hermetic: black is optional here, and letting it + # run would make the scratch output depend on whether it is installed. # Run via `python -m graphblas.core.infixmethods` comparisons = { "lt": "lt", @@ -426,11 +430,17 @@ def _main(): " setattr(VectorIndexExpr, name, val)\n" " setattr(MatrixIndexExpr, name, val)\n" ) + import functools from pathlib import Path from .utils import _autogenerate_code - _autogenerate_code(Path(__file__), "\n".join(lines)) + if not callblack: + _autogenerate_code = functools.partial(_autogenerate_code, callblack=False) + + if base_dir is None: + base_dir = Path(__file__).parent + _autogenerate_code(base_dir / "infixmethods.py", "\n".join(lines)) if __name__ == "__main__": diff --git a/graphblas/core/mask.py b/graphblas/core/mask.py index 4a8412a60..f812e2119 100644 --- a/graphblas/core/mask.py +++ b/graphblas/core/mask.py @@ -75,8 +75,6 @@ def new(self, dtype=None, *, complement=False, mask=None, name=None, **opts): val(self, **opts) << True return val - from .base import _check_mask - mask = _check_mask(mask) d = _COMPLEMENT_MASKS if complement else _COMBINE_MASKS func = d[type(self), type(mask)] @@ -95,8 +93,6 @@ def __and__(self, other, **opts): This uses faster recipes than the above for all combinations of input mask types, and aims to be memory efficient when operating on complemented masks. """ - from .base import _check_mask - other = _check_mask(other) complement = self.complement or other.complement d = _COMPLEMENT_MASKS if complement else _COMBINE_MASKS @@ -121,8 +117,6 @@ def __or__(self, other, **opts): This uses faster recipes than the above for all combinations of input mask types, and aims to be memory efficient when operating on complemented masks. """ - from .base import _check_mask - other = _check_mask(other) func = _MASK_OR[type(self), type(other)] return func(self, other, opts) @@ -202,6 +196,31 @@ def _name_html(self): return f"~{self.parent._name_html}.V" +def _check_mask(mask, output=None, strict_kind=False): + if not isinstance(mask, Mask): + # Convert bool objects to value masks + if utils.output_type(mask).__name__ in {"Vector", "Matrix"}: + if mask.dtype != BOOL: + raise TypeError( + f"Mask must be boolean objects (got {mask.dtype}) " + "or indicate values (M.V) or structure (M.S)" + ) + mask = mask.V # auto-compute (will raise if disabled) + else: + raise TypeError(f"Invalid mask: {type(mask)}") + if output is not None: + if output.ndim == 1 and mask.parent.ndim != 1: + raise TypeError(f"Mask object must be type Vector; got {type(mask.parent)}") + # A full-tensor op (ewise, mxm, apply, extract, ...) into a Matrix needs + # a Matrix mask. Assignment is exempt (`strict_kind` stays False for it): + # a Vector mask on a Matrix row/column assign is valid and is validated + # separately in Matrix.__setitem__. Without this, a Vector mask on a + # full-Matrix op leaked a raw cffi "struct GB_Matrix_opaque" error. + if strict_kind and output.ndim == 2 and mask.parent.ndim != 2: + raise TypeError(f"Mask object must be type Matrix; got {type(mask.parent)}") + return mask + + # Recipes to combine two masks. # Legend: # A: any diff --git a/graphblas/core/matrix.py b/graphblas/core/matrix.py index 698bf7b4f..526833eda 100644 --- a/graphblas/core/matrix.py +++ b/graphblas/core/matrix.py @@ -5,12 +5,20 @@ from .. import backend, binary, monoid, select, semiring from ..dtypes import _INDEX, FP64, INT64, lookup_dtype, unify -from ..exceptions import DimensionMismatch, InvalidValue, NoValue, check_status +from ..exceptions import ( + DimensionMismatch, + GrB_NO_VALUE, + InvalidValue, + NoValue, + check_status, + check_status_carg, +) from . import _supports_udfs, automethods, ffi, lib, utils -from .base import BaseExpression, BaseType, _check_mask, call +from .base import BaseExpression, BaseType, _is_recording, call from .descriptor import lookup as descriptor_lookup +from .dtypes import _raise_dtype_or_arraylike from .expr import _ALL_INDICES, AmbiguousAssignOrExtract, IndexerResolver, InfixExprBase, Updater -from .mask import Mask, StructuralMask, ValueMask +from .mask import Mask, StructuralMask, ValueMask, _check_mask from .operator import ( UNKNOWN_OPCLASS, _get_typed_op_from_exprs, @@ -180,7 +188,7 @@ class Matrix(BaseType): """ - __slots__ = "_nrows", "_ncols", "_parent", "ss" + __slots__ = "_nrows", "_ncols", "_parent" ndim = 2 _is_transposed = False _name_counter = itertools.count() @@ -189,7 +197,10 @@ class Matrix(BaseType): def __new__(cls, dtype=FP64, nrows=0, ncols=0, *, name=None): self = object.__new__(cls) - self.dtype = lookup_dtype(dtype) + try: + self.dtype = lookup_dtype(dtype) + except (ValueError, TypeError) as exc: + _raise_dtype_or_arraylike("Matrix", dtype, exc) nrows = _as_scalar(nrows, _INDEX, is_cscalar=True) ncols = _as_scalar(ncols, _INDEX, is_cscalar=True) self.name = f"M_{next(Matrix._name_counter)}" if name is None else name @@ -198,8 +209,6 @@ def __new__(cls, dtype=FP64, nrows=0, ncols=0, *, name=None): self._nrows = nrows.value self._ncols = ncols.value self._parent = None - if backend == "suitesparse": - self.ss = ss(self) return self @classmethod @@ -211,8 +220,6 @@ def _from_obj(cls, gb_obj, dtype, nrows, ncols, *, parent=None, name=None): self._nrows = nrows self._ncols = ncols self._parent = parent - if backend == "suitesparse": - self.ss = ss(self) return self def __del__(self): @@ -337,6 +344,50 @@ def __setitem__(self, keys, expr, **opts): M[0, 0:3] = 17 """ + # Fast path for `A[i, j] = scalar`: a plain (row, col) integer pair and + # an exact-fit Python scalar, with no mask/accum/opts, a non-UDT dtype, + # and no active Recorder. Mirrors Updater -> _assign_element for a single + # element while skipping the resolver, Updater, and Scalar objects. Only + # int/float/bool/complex are taken here so dtype inference and cffi + # coercion match _assign_element exactly; everything else (slices, fancy + # indexing, numpy or Scalar values, `A(mask)[i, j] << x`) falls back to + # the full assign path, leaving mask/accum and coercion unchanged. + if ( + not opts + and type(keys) is tuple + and len(keys) == 2 + and type(expr) in (int, float, bool, complex) + and not self.dtype._is_udt + and not _is_recording() + ): + row, col = keys + # Only int and np.integer keys take the fast lane (bools excluded), + # matching parse_index; other __index__ objects fall through to the + # expression path and its canonical errors. + if (type(row) is int or isinstance(row, np.integer)) and ( + type(col) is int or isinstance(col, np.integer) + ): + rowidx = row.__index__() + colidx = col.__index__() + nrows = self._nrows + ncols = self._ncols + if rowidx < 0: + rowidx += nrows + if rowidx < 0 or rowidx >= nrows: + raise IndexError(f"Index out of range: index={row}, size={nrows}") + if colidx < 0: + colidx += ncols + if colidx < 0 or colidx >= ncols: + raise IndexError(f"Index out of range: index={col}, size={ncols}") + vdtype = lookup_dtype(type(expr), expr) + cvalue = ffi_new(f"{vdtype.c_type}*") + cvalue[0] = expr # cffi coercion, identical to the Scalar.value setter + err_code = utils.libget(f"GrB_Matrix_setElement_{vdtype.name}")( + self.gb_obj[0], cvalue[0], rowidx, colidx + ) + if err_code: + check_status_carg(err_code, "Matrix", self.gb_obj[0]) + return Updater(self, opts=opts)[keys] = expr def __contains__(self, index): @@ -349,6 +400,46 @@ def __contains__(self, index): (10, 15) in M """ + # Fast path for a plain (row, col) integer pair: probe with + # GrB_Matrix_extractElement directly instead of building an extract + # expression and Scalar. An out-of-range index falls through to the + # expression path so it raises the same IndexError as the slow path. + # Fall back for a UDT dtype and an active Recorder (so the call is + # recorded), mirroring Matrix.get. TransposedMatrix reuses this method. + if ( + type(index) is tuple + and len(index) == 2 + and not self.dtype._is_udt + and not _is_recording() + ): + row, col = index + # Only int and np.integer take the fast lane (bools excluded), + # matching parse_index; other __index__ objects fall through to the + # expression path and its canonical errors. + if (type(row) is int or isinstance(row, np.integer)) and ( + type(col) is int or isinstance(col, np.integer) + ): + rowidx = row.__index__() + colidx = col.__index__() + nrows = self._nrows + ncols = self._ncols + if rowidx < 0: + rowidx += nrows + if colidx < 0: + colidx += ncols + if 0 <= rowidx < nrows and 0 <= colidx < ncols: + if self._is_transposed: + rowidx, colidx = colidx, rowidx + dtype = self.dtype + res = ffi_new(f"{dtype.c_type}*") + err_code = utils.libget(f"GrB_Matrix_extractElement_{dtype.name}")( + res, self.gb_obj[0], rowidx, colidx + ) + if err_code: + if err_code == GrB_NO_VALUE: + return False + check_status_carg(err_code, "Matrix", self.gb_obj[0]) + return True extractor = self[index] if not extractor._is_scalar: raise TypeError( @@ -805,6 +896,46 @@ def get(self, row, col, default=None): Python scalar """ + # Fast path for plain integer indices: call GrB_Matrix_extractElement + # directly instead of building an extract expression, which costs ~10x + # more than the C call for single-element access. Fall back when a + # Recorder is active (so the call is recorded) and for UDTs (whose + # values need numpy-based conversion in Scalar.value). + # Only int and np.integer take the fast lane (bools excluded), matching + # parse_index; other __index__ objects fall through to the expression + # path and its canonical errors. + if ( + (type(row) is int or isinstance(row, np.integer)) + and (type(col) is int or isinstance(col, np.integer)) + and not self.dtype._is_udt + and not _is_recording() + ): + rowidx = row.__index__() + colidx = col.__index__() + nrows = self._nrows + ncols = self._ncols + if rowidx < 0: + rowidx += nrows + if rowidx < 0 or rowidx >= nrows: + raise IndexError(f"Index out of range: index={row}, size={nrows}") + if colidx < 0: + colidx += ncols + if colidx < 0 or colidx >= ncols: + raise IndexError(f"Index out of range: index={col}, size={ncols}") + if self._is_transposed: + # TransposedMatrix reuses this method; gb_obj is the + # untransposed parent, so extract the mirrored element. + rowidx, colidx = colidx, rowidx + dtype = self.dtype + res = ffi_new(f"{dtype.c_type}*") + err_code = utils.libget(f"GrB_Matrix_extractElement_{dtype.name}")( + res, self.gb_obj[0], rowidx, colidx + ) + if err_code: + if err_code == GrB_NO_VALUE: + return default + check_status_carg(err_code, "Matrix", self.gb_obj[0]) + return res[0] expr = self[row, col] if expr._is_scalar: rv = expr.new().value @@ -829,6 +960,12 @@ def from_coo( ): """Create a new Matrix from row and column indices and values. + .. warning:: + When ``nrows`` or ``ncols`` is omitted, the shape is inferred from + the largest row or column index, so trailing all-empty rows or + columns are dropped. Pass ``nrows`` and ``ncols`` explicitly to pin + the shape. + Parameters ---------- rows : list or np.ndarray @@ -1590,10 +1727,7 @@ def from_dicts( else: # If we know the dtype, then using `np.fromiter` is much faster dtype = lookup_dtype(dtype) - if dtype.np_type.subdtype is not None and np.__version__[:5] in {"1.21.", "1.22."}: - values, dtype = values_to_numpy_buffer(list(iter_values), dtype) # FLAKY COVERAGE - else: - values = np.fromiter(iter_values, dtype.np_type) + values = np.fromiter(iter_values, dtype.np_type) return getattr(cls, methodname)( *args, indptr, col_indices, values, dtype, nrows=nrows, ncols=ncols, name=name ) @@ -3282,6 +3416,13 @@ def _prep_for_assign(self, resolved_indexes, value, mask, is_submask, replace, o within=method_name, extra_message=extra_message, ) + if mask is not None and mask.parent.ndim != 2: + # Matrix value, Vector mask, Matrix index + # C(m)[I, J] << A + # C[I, J](m) << A + # This also catches whole-object updates such as `C(m) << A` + # and `C.dup(mask=m)`, which assign with `C(m)[...] = A`. + raise TypeError("Unable to use Vector mask on Matrix assignment to a Matrix") if is_submask: # C[I, J](M) << A expr_repr = ( @@ -3532,10 +3673,20 @@ def _delete_element(self, resolved_indexes): if backend == "suitesparse": - Matrix.ss = class_property(Matrix.ss, ss) + # `.ss` is built lazily per access rather than stored on the instance: a + # stored `ss(self)` holds `_parent` back to this Matrix (and so does its + # config), forming a reference cycle that keeps the object (and its C-side + # GrB buffer) alive until the cyclic gc runs instead of dying by refcount. + # See gh-559. Class-level `Matrix.ss` still resolves to the `ss` class so + # the `import_*` classmethods keep working. + def _ss(self): + return ss(self) + + _ss.__name__ = _ss.__qualname__ = "ss" + Matrix.ss = class_property(property(_ss), ss) else: Matrix.ss = class_property( - Matrix.ss, 'ss attribute is only available with "suitesparse" backend', exceptional=True + property(), 'ss attribute is only available with "suitesparse" backend', exceptional=True ) diff --git a/graphblas/core/operator/agg.py b/graphblas/core/operator/agg.py index 7afeb9e46..f1100c32e 100644 --- a/graphblas/core/operator/agg.py +++ b/graphblas/core/operator/agg.py @@ -29,6 +29,27 @@ def _get_types(ops, initdtype): class Aggregator: + """A reduction operator that collapses the values of a Matrix or Vector. + + An Aggregator is used with the ``reduce`` family of methods: ``Vector.reduce``, + ``Matrix.reduce_rowwise``, ``Matrix.reduce_columnwise``, and ``Matrix.reduce_scalar``. + Some aggregators return a summary value (``sum``, ``mean``, ``max``); others return a + position (``ss.argmin``, ``ss.argmax``, ``ss.first_index``). + + Built-in aggregators live in the ``graphblas.agg`` namespace, such as ``agg.sum``, + ``agg.mean``, and ``agg.count``. The position aggregators are SuiteSparse-specific and + live under ``agg.ss``; the bare ``agg.argmin`` spellings are deprecated. Unlike the + other operators, an Aggregator is not a single GraphBLAS object; many are built from a + monoid or semiring plus an optional finalize step, composed per dtype on first use. + + Examples + -------- + >>> import graphblas as gb + >>> v = gb.Vector.from_coo([0, 1, 2], [1, 2, 3]) + >>> int(v.reduce(gb.agg.sum).new()) + 6 + """ + opclass = "Aggregator" def __init__( @@ -755,4 +776,6 @@ def _first_last_index(agg, updater, expr, opts, *, in_composite, semiring): agg.Aggregator = Aggregator agg.TypedAggregator = TypedAggregator -from .utils import get_typed_op # noqa: E402 isort:skip +from .utils import _register_aggregator_types, get_typed_op # noqa: E402 isort:skip + +_register_aggregator_types(Aggregator, TypedAggregator) diff --git a/graphblas/core/operator/base.py b/graphblas/core/operator/base.py index bf685ba5c..6d4ec6554 100644 --- a/graphblas/core/operator/base.py +++ b/graphblas/core/operator/base.py @@ -3,10 +3,13 @@ from operator import getitem from types import BuiltinFunctionType, ModuleType +import numpy as np + from ... import _STANDARD_OPERATOR_NAMES, backend, op from ...dtypes import BOOL, INT8, UINT64, _supports_complex, lookup_dtype from ...exceptions import UdfParseError, check_status_carg from .. import _has_numba, _supports_udfs, ffi, lib +from ..dtypes import _sample_values from ..expr import InfixExprBase from ..utils import output_type @@ -109,6 +112,39 @@ def _bool_to_int8(dtype): return INT8 if dtype == BOOL else dtype +def _validate_ret_dtype(ret_dtype, opclass, *, is_udt, parameterized): + """Normalize a user-supplied ``ret_dtype`` to a DataType, or raise. + + ``ret_dtype`` names the operator's output type outright instead of letting + it be inferred from what the UDF returns. Inference can only name a type it + can see, which is why this is limited to the UDT path: the builtin path + derives its output from Numba's typing of each sample input, and forcing a + single type across all of them would silently recast results. + """ + if ret_dtype is None: + return None + if not is_udt: + raise ValueError( + f"{opclass}: ret_dtype requires is_udt=True. The return type for builtin " + f"dtypes comes from compiling the function for each input type, so a single " + f"fixed type cannot describe it." + ) + if parameterized: + raise ValueError( + f"{opclass}: ret_dtype does not work with parameterized=True. " + f"A parameterized operator builds and registers its function when called, " + f"and that inner registration does not accept a return dtype; register the " + f"built function without parameterized=True to declare one." + ) + try: + return lookup_dtype(ret_dtype) + except (ValueError, TypeError) as exc: + raise ValueError( + f"{opclass}: ret_dtype={ret_dtype!r} is not a recognized dtype. " + f"Pass a DataType, a numpy dtype, or a name such as 'FP64'." + ) from exc + + class OpPath: def __init__(self, parent, name): self._parent = parent @@ -189,7 +225,7 @@ def _finalize_udt_op(parent_op, dtype, dtype2, ret_type, wrapper, wrapper_sig, t from ``typed_user_cls.opclass``. ``dtype2`` is ``None`` for unary ops; the rest pass both. Returns the cached ``TypedUser*Op``. """ - wrapper = numba.cfunc(wrapper_sig, nopython=True)(wrapper) + wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")(wrapper) c_typename = _GB_OBJ_C_TYPENAME[typed_user_cls.opclass] error_label = c_typename.removeprefix("GrB_").removeprefix("GxB_") gb_obj = ffi.new(f"{c_typename}*") @@ -267,6 +303,16 @@ def _summarize_numba_typing_error(exc): return line return "Numba could not compile the function for these input types" + def _udt_ret_type(parent_op, numba_ret_type, *dtypes): + """Return the operator's declared ``ret_dtype``, else infer one from the UDF. + + Inference can only name a type that is already an operand, so an output + UDT that appears nowhere in the inputs is unreachable without this. + """ + if (ret_dtype := parent_op._ret_dtype) is not None: + return ret_dtype + return _resolve_udt_return_type(numba_ret_type, *dtypes) + def _resolve_udt_return_type(numba_ret_type, *dtypes): """Resolve a Numba return type to a DataType, matching Tuple returns to an input UDT. @@ -329,17 +375,80 @@ def _leaves(d): f"shape {shape}. Return a numpy array (e.g., ``np.array(...)``) or a " f"scalar; tuple returns are only matched to record UDTs." ) + elif isinstance(numba_ret_type, numba.core.types.Array): + # A UDF over an array UDT may build its result (``x + y``) instead + # of returning an operand. Numba types that as a plain Array, which + # ``lookup_dtype`` doesn't recognize, so match it back to an array + # UDT input by base element type and dimensionality. + candidates = [ + d + for d in dtypes + if d._is_udt + and d.np_type.subdtype is not None + and d.numba_type.dtype == numba_ret_type.dtype + and len(d.numba_type.shape) == numba_ret_type.ndim + ] + # An Array type carries ``ndim`` but not its extents, so operands + # that differ only in length are indistinguishable here. Guessing + # would hand SuiteSparse an element of the wrong size, so say so. + # Compare the UDTs rather than their shapes: a flat ``FP64[2, 3]`` + # and a layered ``FP64[3][2]`` are separate DataTypes with separate + # GrB_Type handles, yet Numba collapses both to the same shape. + unique = [] + for d in candidates: + if not any(d is seen for seen in unique): + unique.append(d) + if len(unique) > 1: + raise UdfParseError( + f"UDT UDF returned {numba_ret_type!r}, which matches more than one " + f"input array UDT ({', '.join(str(d) for d in unique)}). " + f"Return one of the operands, or make the operands the same type." + ) + if unique: + return unique[0] + # An array UDT went in and an array came out, but not one that + # fits: name the mismatch rather than fall through to the generic + # "unsupported type", whose advice the user already followed. + array_inputs = [d for d in dtypes if d._is_udt and d.np_type.subdtype is not None] + if array_inputs: + d = array_inputs[0] + nested = d.numba_type + raise UdfParseError( + f"UDT UDF returned {numba_ret_type!r}, which matches no input array " + f"UDT: {d} elements are {nested.dtype} with shape {nested.shape}. " + f"Return an array of that dtype and rank, or one of the operands." + ) raise UdfParseError( f"UDT UDF returned an unsupported type {numba_ret_type!r}. " f"Return a scalar, a tuple matching a record UDT's fields, or a numpy array " f"matching an array UDT's shape." ) + def _array_udt_view(dtype): + """Return ``(base_element_numba_type, shape)`` for an array UDT. + + Array UDTs are addressed as a ``carray`` over their base elements, in + the UDT's declared shape, rather than as Numba's ``NestedArray``. Numba + models a ``NestedArray`` *value* as an array descriptor (data pointer, + shape, strides, ...), so loading or storing one through a ``CPointer`` + moves the descriptor rather than the element payload, corrupting + whatever follows it (and overrunning the element outright once the + descriptor is the wider of the two). + + Read from ``numba_type`` rather than ``np_type.subdtype`` so this + agrees with the type the UDF was compiled against: numpy keeps nested + subarray dtypes layered, e.g. ``FP64[5]`` inside ``[6]`` stays + ``(dtype((' len(expected) and shape[0] == 1: + shape = shape[1:] + try: + return np.broadcast_shapes(shape, expected) == expected + except ValueError: + return False + + def _check_array_udf_shape(numba_func, return_type, operands): + """Reject a UDF whose built array cannot fill the array UDT's element.""" + probed = _run_udf_probe(numba_func, operands) + if probed is None: + return + expected = return_type.numba_type.shape + shape = getattr(probed[0], "shape", None) + if shape is not None and not _fits_by_broadcast(shape, expected): + raise UdfParseError( + f"UDT UDF returned an array of shape {tuple(shape)} when run on sample " + f"values, but {return_type} elements are {tuple(expected)}. Return an " + f"array whose shape matches or broadcasts to that, or one of the " + f"operands. {_SHAPE_HINT}" + ) + + def _check_record_udf_leaf_shapes(numba_func, return_type, operands): + """Reject a record return whose array-typed leaf cannot fill its field. + + The array-leaf half of :func:`_check_array_udf_shape`. The wrapper + slice-assigns those leaves, so a wrong extent raises inside the cfunc + and abandons the write part-way: every leaf after it keeps whatever + SuiteSparse had in the buffer, scalar leaves included. + + Only records that actually have an array leaf are probed, so the + common record UDF still registers without running the user's code. + """ + from .udt_utils import _iter_record_leaves + + leaves = [(py, d) for py, _c, d in _iter_record_leaves(return_type.np_type)] + if not any(d.subdtype is not None for _py, d in leaves): + return + probed = _run_udf_probe(numba_func, operands) + if probed is None: + return + result = probed[0] + if not isinstance(result, tuple) or len(result) != len(leaves): + return + for (path, leaf_dtype), value in zip(leaves, result, strict=True): + if leaf_dtype.subdtype is None: + continue + expected = leaf_dtype.subdtype[1] + shape = getattr(value, "shape", None) + if shape is not None and not _fits_by_broadcast(shape, expected): + raise UdfParseError( + f"UDT UDF returned an array of shape {tuple(shape)} for field " + f"{path} of {return_type} when run on sample values, which holds " + f"{tuple(expected)} there. Return an array whose shape matches or " + f"broadcasts to that. {_SHAPE_HINT}" + ) + def _get_udt_wrapper( numba_func, return_type, dtype, dtype2=None, *, include_indexes=False, numba_ret_type=None ): @@ -440,15 +723,22 @@ def _get_udt_wrapper( zsetup, zptr_type, zkind, zinfo = _output_handler(return_type, numba_ret_type) xsetup, xderef, xptr_type = _input_operand(dtype, "x") wrapper_args = [zptr_type, xptr_type] + probe_operands = [dtype] if include_indexes: wrapper_args.extend([UINT64.numba_type, UINT64.numba_type]) + probe_operands.extend([UINT64, UINT64]) ysetup, yderef_expr, yarg = "", "", "" if dtype2 is not None: ysetup, yderef, yptr_type = _input_operand(dtype2, "y") wrapper_args.append(yptr_type) + probe_operands.append(dtype2) yarg = ", y_ptr" yderef_expr = f", {yderef}" wrapper_sig = nt.void(*wrapper_args) + if zkind == "array_elements": + _check_array_udf_shape(numba_func, return_type, probe_operands) + elif zkind == "record_fields": + _check_record_udf_leaf_shapes(numba_func, return_type, probe_operands) rcidx = ", row, col" if include_indexes else "" signature_line = f"def wrapper(z_ptr, x_ptr{rcidx}{yarg}):" @@ -490,6 +780,16 @@ def _get_udt_wrapper_indexbinary( UINT64.numba_type, tptr_type, ) + # Same registration-time guards as :func:`_get_udt_wrapper`. Without + # them a wrong-shape return is memory-safe but silent: the wrapper's + # ``z[:] =`` raises inside the cfunc, Numba prints and swallows it, + # and the element is left as SuiteSparse found it. + if zkind in ("array_elements", "record_fields"): + probe_operands = [dtype, UINT64, UINT64, dtype2, UINT64, UINT64, dtype2] + if zkind == "array_elements": + _check_array_udf_shape(numba_func, return_type, probe_operands) + else: + _check_record_udf_leaf_shapes(numba_func, return_type, probe_operands) signature_line = "def wrapper(z_ptr, x_ptr, ix, jx, y_ptr, iy, jy, t_ptr):" body_setup = f"{zsetup}{xsetup}{ysetup}{tsetup}" @@ -534,16 +834,13 @@ class TypedOpBase: "gb_name", "_type2", "_jit_c_info", - "_owns_gb_obj_inst", + "_gb_obj_owner", "__weakref__", ) # Subclasses whose ``gb_obj`` was allocated via ``GrB__new`` / # ``GxB__new`` (TypedUser*Op, _BoundIndexBinaryOp) override this so # ``__del__`` frees the SuiteSparse handle. Built-in typed ops point at # SuiteSparse's permanent built-in singletons and must never free. - # Specific instances can override via ``_owns_gb_obj_inst`` (set by - # the constructor); ``SelectOp._from_indexunary`` aliases an existing - # ``GrB_IndexUnaryOp`` and must clear ownership to avoid a double free. _owns_gb_obj = False def __init__(self, parent, name, type_, return_type, gb_obj, gb_name, dtype2=None): @@ -558,10 +855,12 @@ def __init__(self, parent, name, type_, return_type, gb_obj, gb_name, dtype2=Non # for this typed op; ``None`` for built-in ops and for UDT ops with # no JIT path. self._jit_c_info = None - # Per-instance ownership override; defaults to the class attribute. - # ``SelectOp._from_indexunary`` flips this to ``False`` on aliasing - # TypedUserSelectOps so only the IndexUnaryOp frees the handle. - self._owns_gb_obj_inst = type(self)._owns_gb_obj + # Set when ``gb_obj`` is borrowed from another typed op rather than + # allocated here (``SelectOp._from_indexunary``). Storing the owner + # both suppresses our free and keeps the handle alive as long as we + # point at it; disclaiming ownership without naming an owner would + # leave this op holding a dangling handle. + self._gb_obj_owner = None @property def jit_c_name(self): @@ -594,11 +893,14 @@ def __reduce__(self): def __del__(self): # Free the SuiteSparse handle we allocated. Built-in typed ops alias - # SuiteSparse's permanent built-in singletons and must never free, so - # gate on the per-instance owns flag (defaults to the class - # attribute; the alias case overrides to False). Mirrors the - # ``Matrix.__del__`` / ``Vector.__del__`` pattern. - if not getattr(self, "_owns_gb_obj_inst", False): + # SuiteSparse's permanent built-in singletons and must never free. + # Mirrors the ``Matrix.__del__`` / ``Vector.__del__`` pattern. + if not type(self)._owns_gb_obj: + return + # A borrowed handle belongs to ``_gb_obj_owner``, which we keep alive. + # ``getattr`` guards the case where ``__init__`` raised before the slot + # was set. + if getattr(self, "_gb_obj_owner", None) is not None: return gb_obj = getattr(self, "gb_obj", None) if gb_obj is None or lib is None or ffi is None: @@ -718,6 +1020,23 @@ def __init__(self, name, *, anonymous=False): def __repr__(self): return f"{self._modname}.{self.name}" + def _build_deferred(self, type_): + """Build a typed op on demand for ``type_``, or return None if this op + has no deferred builds. Overridden where ``.types`` is populated before + the per-dtype numba compilation (see ``BinaryOp``). + """ + return + + def _materialize_deferred(self): + """Build every typed op that ``.types`` advertises but that has not been + compiled yet. A no-op unless ``_build_deferred`` defers builds; call it + before code that iterates ``_typed_ops`` directly (e.g. building a + Semiring over this op). + """ + for type_ in list(self.types): + if type_ not in self._typed_ops: + self._build_deferred(type_) + def __getitem__(self, type_): if type(type_) is tuple: from .utils import get_typed_op @@ -729,6 +1048,9 @@ def __getitem__(self, type_): if not self._is_udt: type_ = lookup_dtype(type_) if type_ not in self._typed_ops: + op = self._build_deferred(type_) + if op is not None: + return op if self._udt_types is None: if self.is_positional: return self._typed_ops[UINT64] @@ -899,21 +1221,25 @@ def _deserialize(cls, name, *args): return cls.register_new(name, *args) @classmethod - def _deserialize_udf(cls, name, orig_func, is_udt): + def _deserialize_udf(cls, name, orig_func, is_udt, ret_dtype=None): """Re-register a named UDF on unpickle, or reuse if already present. Shared by the five UDF-capable subclasses (UnaryOp, BinaryOp, IndexUnaryOp, SelectOp, IndexBinaryOp), all of which use the - default ``__reduce__`` below. + default ``__reduce__`` below. ``ret_dtype`` is passed only when set: + SelectOp shares this path and takes no ret_dtype, and pickles written + before ret_dtype existed carry a 3-tuple. """ if (rv := cls._find(name)) is not None: return rv - return cls.register_new(name, orig_func, is_udt=is_udt) + kwargs = {} if ret_dtype is None else {"ret_dtype": ret_dtype} + return cls.register_new(name, orig_func, is_udt=is_udt, **kwargs) @classmethod - def _deserialize_anon_udf(cls, func, name, is_udt): + def _deserialize_anon_udf(cls, func, name, is_udt, ret_dtype=None): """Re-register an anonymous UDF on unpickle.""" - return cls.register_anonymous(func, name, is_udt=is_udt) + kwargs = {} if ret_dtype is None else {"ret_dtype": ret_dtype} + return cls.register_anonymous(func, name, is_udt=is_udt, **kwargs) def __reduce__(self): """Default ``__reduce__`` for UDF-capable subclasses. @@ -926,10 +1252,16 @@ def __reduce__(self): if self._anonymous: if hasattr(self.orig_func, "_parameterized_info"): return (_deserialize_parameterized, self.orig_func._parameterized_info) - return (type(self)._deserialize_anon_udf, (self.orig_func, self.name, self._is_udt)) + return ( + type(self)._deserialize_anon_udf, + (self.orig_func, self.name, self._is_udt, getattr(self, "_ret_dtype", None)), + ) if (name := f"{self._modname}.{self.name}") in _STANDARD_OPERATOR_NAMES: return name - return (type(self)._deserialize_udf, (self.name, self.orig_func, self._is_udt)) + return ( + type(self)._deserialize_udf, + (self.name, self.orig_func, self._is_udt, getattr(self, "_ret_dtype", None)), + ) @classmethod def _check_supports_udf(cls, method_name): diff --git a/graphblas/core/operator/binary.py b/graphblas/core/operator/binary.py index 97e0a9e70..fffbf89d3 100644 --- a/graphblas/core/operator/binary.py +++ b/graphblas/core/operator/binary.py @@ -33,6 +33,7 @@ TypedOpBase, _call_op, _hasop, + _validate_ret_dtype, ) # Imported unconditionally (plain dict, no numba): ``_compile_udt`` consults it @@ -47,8 +48,28 @@ _compile_udf_for_udt, _finalize_udt_op, _get_udt_wrapper, - _resolve_udt_return_type, + _udt_ret_type, ) + + try: + # Typing-only return-type inference for the module-level built-in UDFs: + # learn return types without lowering or codegen, so ``.types`` can be + # populated up front while the njit compile is deferred to first use. + # Numba's public API offers no way to ask for a return type without + # also compiling for it, so this reaches into ``numba.core``. That is + # internal, and we support numba back to 0.57, so treat every part of + # it as breakable: a failed import or any surprise from the call leaves + # ``_infer_ret_types_typing_only`` returning None and ``_build`` on the + # eager full-compile loop, which is what it did before this existed. + from numba.core import compiler as _numba_compiler + from numba.core import typed_passes as _numba_typed_passes + from numba.core.registry import cpu_target as _numba_cpu_target + + _HAS_TYPING_ONLY = True + except Exception: # pragma: no cover - numba internals moved + _HAS_TYPING_ONLY = False +else: + _HAS_TYPING_ONLY = False if _supports_complex: from ...dtypes import FC32, FC64 @@ -202,13 +223,13 @@ def monoid(self): @property def commutes_to(self): commutes_to = self.parent.commutes_to - if commutes_to is not None and (self.type in commutes_to._typed_ops or self.type._is_udt): + if commutes_to is not None and (self.type in commutes_to.types or self.type._is_udt): return commutes_to[self.type] @property def _semiring_commutes_to(self): commutes_to = self.parent._semiring_commutes_to - if commutes_to is not None and (self.type in commutes_to._typed_ops or self.type._is_udt): + if commutes_to is not None and (self.type in commutes_to.types or self.type._is_udt): return commutes_to[self.type] @property @@ -388,6 +409,125 @@ def _pair_dtype(op, dtype, dtype2): return op[INT64] +def _adjust_ret_type(type_, ret_type, return_types): + """Downcast a UDF's inferred return type toward the input type when that is + the intent (INT->INT, FP->FP, FC->FC, and the UINT64/BOOL special cases). + + ``return_types`` holds the results already decided for earlier sample dtypes; + the UINT64 and BOOL rules consult INT64 and INT8, which are decided first. + Shared by the eager compile loop and the typing-only inference path so both + produce the same ``.types``. + """ + if ret_type != type_ and ( + ("INT" in ret_type.name and "INT" in type_.name) + or ("FP" in ret_type.name and "FP" in type_.name) + or ("FC" in ret_type.name and "FC" in type_.name) + or (type_ == UINT64 and ret_type == FP64 and return_types.get(INT64) == INT64) + ): + # This is what users want most of the time, but we can't make a perfect + # rule. There should be a way for users to be explicit. + return type_ + if type_ == BOOL and ret_type == INT64 and return_types.get(INT8) == INT8: + return INT8 + return ret_type + + +def _finalize_typed_binaryop(parent, binary_udf, name, type_, ret_type): + """Compile the cfunc wrapper for one dtype and register the ``GrB_BinaryOp``. + + Shared by the eager build loop and the deferred per-dtype build. ``ret_type`` + is already resolved (heuristic applied); this lowers ``binary_udf`` for the + dtype (via the cfunc), wires up the GraphBLAS op, and records the typed op. + For deferred built-ins this is where that dtype pays its compilation cost. + """ + nt = numba.types + input_type = _bool_to_int8(type_) + return_type = _bool_to_int8(ret_type) + + # Build wrapper because GraphBLAS wants pointers and void return + wrapper_sig = nt.void( + nt.CPointer(return_type.numba_type), + nt.CPointer(input_type.numba_type), + nt.CPointer(input_type.numba_type), + ) + + if type_ == BOOL: + if ret_type == BOOL: + + def binary_wrapper(z, x, y): # pragma: no cover (numba) + z[0] = bool(binary_udf(bool(x[0]), bool(y[0]))) + + else: + + def binary_wrapper(z, x, y): # pragma: no cover (numba) + z[0] = binary_udf(bool(x[0]), bool(y[0])) + + elif ret_type == BOOL: + + def binary_wrapper(z, x, y): # pragma: no cover (numba) + z[0] = bool(binary_udf(x[0], y[0])) + + else: + + def binary_wrapper(z, x, y): # pragma: no cover (numba) + z[0] = binary_udf(x[0], y[0]) + + binary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")(binary_wrapper) + new_binary = ffi_new("GrB_BinaryOp*") + check_status_carg( + lib.GrB_BinaryOp_new( + new_binary, + binary_wrapper.cffi, + ret_type.gb_obj, + type_.gb_obj, + type_.gb_obj, + ), + "BinaryOp", + new_binary[0], + ) + op = TypedUserBinaryOp(parent, name, type_, ret_type, new_binary[0]) + parent._add(op) + return op + + +def _infer_ret_types_typing_only(func): + """Infer ``{DataType: return DataType}`` for the sample dtypes without + lowering ``func``, or return None on any failure so the caller falls back to + the eager full-compile loop. + + Used only for the module-level built-in UDFs, whose typing and lowering are + known to agree (verified: typing-only reproduces their live ``.types`` + exactly). Typing is more permissive than lowering in general, so this must + not be used for arbitrary user funcs, where a dtype that types but fails to + lower would be wrongly reported as supported. + """ + if not _HAS_TYPING_ONLY: + return None + try: + typingctx = _numba_cpu_target.typing_context + targetctx = _numba_cpu_target.target_context + typingctx.refresh() + targetctx.refresh() + return_types = {} + for type_ in _sample_values: + try: + interp = _numba_compiler.run_frontend(func) + result = _numba_typed_passes.type_inference_stage( + typingctx, targetctx, interp, [type_.numba_type, type_.numba_type], None + ) + numba_ret_type = result.return_type + except numba.TypingError: + # This dtype does not type-check; skip it, matching the eager + # loop's ``except numba.TypingError: continue``. + continue + ret_type = _adjust_ret_type(type_, lookup_dtype(numba_ret_type), return_types) + return_types[type_] = ret_type + except Exception: # pragma: no cover - unexpected numba result shape/API + return None + else: + return return_types + + if _has_numba: from .udt_utils import ( _compile_codegen, @@ -414,6 +554,8 @@ class BinaryOp(OpBase): "_is_udt", "_numba_func", "_custom_dtype", + "_defer_builds", + "_ret_dtype", ) _module = binary _modname = "binary" @@ -519,85 +661,67 @@ class BinaryOp(OpBase): } @classmethod - def _build(cls, name, func, *, is_udt=False, anonymous=False): + def _build(cls, name, func, *, is_udt=False, anonymous=False, cache=False, ret_dtype=None): if not isinstance(func, FunctionType): raise TypeError(f"UDF argument must be a function, not {type(func)}") if name is None: name = getattr(func, "__name__", "") + # This rejects ret_dtype unless is_udt, which keeps it disjoint from the + # deferred builtin path below: that one only runs under ``not is_udt``, + # so an op can never be both deferred and carrying a declared return + # type. ``_build_deferred`` reads its ret_type back from ``.types`` and + # never consults ``_ret_dtype``. + ret_dtype = _validate_ret_dtype(ret_dtype, "binary", is_udt=is_udt, parameterized=False) success = False - binary_udf = numba.njit(func) - new_type_obj = cls(name, func, anonymous=anonymous, is_udt=is_udt, numba_func=binary_udf) + # The error model has to be set here as well as on the cfunc wrapper in + # ``_finalize_typed_binaryop``. A Dispatcher keeps one compilation per + # signature, so whichever compile happens first fixes the model for + # that signature, and the eager loop's ``.compile(sig)`` gets there + # before the wrapper does. Miss either one and ``x // 0`` raises + # ZeroDivisionError inside a cfunc, where Numba prints the traceback + # and returns, handing GraphBLAS an element it never wrote. + # + # ``cache`` is only True for the module-level built-in UDFs. Numba can + # only key a cache entry on a stable on-disk source, which lambdas and + # interactively defined functions do not have, so ops registered by + # users stay uncached. + binary_udf = numba.njit(func, error_model="numpy", cache=cache) + new_type_obj = cls( + name, + func, + anonymous=anonymous, + is_udt=is_udt, + numba_func=binary_udf, + ret_dtype=ret_dtype, + ) return_types = {} - nt = numba.types if not is_udt: - for type_ in _sample_values: - sig = (type_.numba_type, type_.numba_type) - try: - binary_udf.compile(sig) - except numba.TypingError: - continue - ret_type = lookup_dtype(binary_udf.overloads[sig].signature.return_type) - if ret_type != type_ and ( - ("INT" in ret_type.name and "INT" in type_.name) - or ("FP" in ret_type.name and "FP" in type_.name) - or ("FC" in ret_type.name and "FC" in type_.name) - or (type_ == UINT64 and ret_type == FP64 and return_types.get(INT64) == INT64) - ): - # Downcast `ret_type` to `type_`. - # This is what users want most of the time, but we can't make a perfect rule. - # There should be a way for users to be explicit. - ret_type = type_ - elif type_ == BOOL and ret_type == INT64 and return_types.get(INT8) == INT8: - ret_type = INT8 - - input_type = _bool_to_int8(type_) - return_type = _bool_to_int8(ret_type) - - # Build wrapper because GraphBLAS wants pointers and void return - wrapper_sig = nt.void( - nt.CPointer(return_type.numba_type), - nt.CPointer(input_type.numba_type), - nt.CPointer(input_type.numba_type), - ) - - if type_ == BOOL: - if ret_type == BOOL: - - def binary_wrapper(z, x, y): # pragma: no cover (numba) - z[0] = bool(binary_udf(bool(x[0]), bool(y[0]))) - - else: - - def binary_wrapper(z, x, y): # pragma: no cover (numba) - z[0] = binary_udf(bool(x[0]), bool(y[0])) - - elif ret_type == BOOL: - - def binary_wrapper(z, x, y): # pragma: no cover (numba) - z[0] = bool(binary_udf(x[0], y[0])) - - else: - - def binary_wrapper(z, x, y): # pragma: no cover (numba) - z[0] = binary_udf(x[0], y[0]) - - binary_wrapper = numba.cfunc(wrapper_sig, nopython=True)(binary_wrapper) - new_binary = ffi_new("GrB_BinaryOp*") - check_status_carg( - lib.GrB_BinaryOp_new( - new_binary, - binary_wrapper.cffi, - ret_type.gb_obj, - type_.gb_obj, - type_.gb_obj, - ), - "BinaryOp", - new_binary[0], - ) - op = TypedUserBinaryOp(new_type_obj, name, type_, ret_type, new_binary[0]) - new_type_obj._add(op) - success = True - return_types[type_] = ret_type + # ``cache=True`` marks the module-level built-in UDFs (floordiv and + # friends). For those, infer return types without lowering so + # ``.types`` is fully populated up front, then defer each dtype's + # njit lowering + cfunc wrapper to first use (see ``_build_deferred``). + # register_new/register_anonymous user funcs keep the eager loop + # below, which validates lowerability for every sample dtype. If + # inference fails for any reason it returns None and we fall back to + # the eager loop too. + inferred = _infer_ret_types_typing_only(func) if cache else None + if inferred is not None: + new_type_obj.types.update(inferred) + return_types.update(inferred) + new_type_obj._defer_builds = True + success = bool(inferred) + else: + for type_ in _sample_values: + sig = (type_.numba_type, type_.numba_type) + try: + binary_udf.compile(sig) + except numba.TypingError: + continue + ret_type = lookup_dtype(binary_udf.overloads[sig].signature.return_type) + ret_type = _adjust_ret_type(type_, ret_type, return_types) + _finalize_typed_binaryop(new_type_obj, binary_udf, name, type_, ret_type) + success = True + return_types[type_] = ret_type if success or is_udt: return new_type_obj raise UdfParseError("Unable to parse function using Numba") @@ -642,7 +766,7 @@ def _compile_udt(self, dtype, dtype2): numba_func, sig, op_kind="binary", op_name=self.name, dtypes=(dtype, dtype2) ) numba_ret_type = numba_func.overloads[sig].signature.return_type - ret_type = _resolve_udt_return_type(numba_ret_type, dtype, dtype2) + ret_type = _udt_ret_type(self, numba_ret_type, dtype, dtype2) binary_wrapper, wrapper_sig = _get_udt_wrapper( numba_func, ret_type, dtype, dtype2, numba_ret_type=numba_ret_type ) @@ -681,7 +805,9 @@ def _compile_udt(self, dtype, dtype2): return op @classmethod - def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=False): + def register_anonymous( + cls, func, name=None, *, parameterized=False, is_udt=False, ret_dtype=None + ): """Register a BinaryOp without registering it in the ``graphblas.binary`` namespace. Because it is not registered in the namespace, the name is optional. @@ -710,6 +836,13 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals Setting ``is_udt=True`` is also helpful when the left and right dtypes need to be different. + ret_dtype : dtype, optional + The dtype the operator returns. Requires ``is_udt=True``. + Without it the return type is inferred from what the function + returns, which can only name a type that is already an input, so + an output UDT that is not an operand needs this. The dtype is + fixed for the operator: it is the same for every input dtype. + Returns ------- BinaryOp or ParameterizedBinaryOp @@ -717,11 +850,22 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals """ cls._check_supports_udf("register_anonymous") if parameterized: + _validate_ret_dtype(ret_dtype, "binary", is_udt=is_udt, parameterized=True) return ParameterizedBinaryOp(name, func, anonymous=True, is_udt=is_udt) - return cls._build(name, func, anonymous=True, is_udt=is_udt) + return cls._build(name, func, anonymous=True, is_udt=is_udt, ret_dtype=ret_dtype) @classmethod - def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=False): + def register_new( + cls, + name, + func, + *, + parameterized=False, + is_udt=False, + lazy=False, + ret_dtype=None, + _cache=False, + ): """Register a new BinaryOp and save it to ``graphblas.binary`` namespace. Parameters @@ -756,6 +900,10 @@ def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=Fal delay compilation and only compile when the operator is used, which is done by setting ``lazy=True``. + ret_dtype : dtype, optional + The dtype the operator returns. Requires ``is_udt=True``. + See :meth:`register_anonymous` for details. + Examples -------- >>> def max_zero(x, y): @@ -779,6 +927,9 @@ def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=Fal """ cls._check_supports_udf("register_new") + # Validate eagerly even for lazy=True, so a bad combination fails at + # the registration site rather than at first attribute touch. + _validate_ret_dtype(ret_dtype, "binary", is_udt=is_udt, parameterized=parameterized) module, funcname = cls._remove_nesting(name) if lazy: module._delayed[funcname] = ( @@ -788,13 +939,16 @@ def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=Fal "func": func, "parameterized": parameterized, "is_udt": is_udt, + "ret_dtype": ret_dtype, + "_cache": _cache, }, ) elif parameterized: + _validate_ret_dtype(ret_dtype, "binary", is_udt=is_udt, parameterized=True) binary_op = ParameterizedBinaryOp(name, func, is_udt=is_udt) setattr(module, funcname, binary_op) else: - binary_op = cls._build(name, func, is_udt=is_udt) + binary_op = cls._build(name, func, is_udt=is_udt, cache=_cache, ret_dtype=ret_dtype) setattr(module, funcname, binary_op) # Also save it to `graphblas.op` if not yet defined opmodule, funcname = cls._remove_nesting(name, module=op, modname="op", strict=False) @@ -847,13 +1001,20 @@ def _initialize(cls): if _supports_udfs: # Add floordiv # cdiv truncates towards 0, while floordiv truncates towards -inf - BinaryOp.register_new("floordiv", _floordiv, lazy=True) # cast to integer - BinaryOp.register_new("rfloordiv", _rfloordiv, lazy=True) # cast to integer + # cache=True persists the numba compilation of these built-in UDFs to + # disk so it is paid once per machine, not once per process. It shaves + # ~20% off a single op's first-touch build and ~45% when several are + # used (the per-process object-link and LLVM-init cost is not cached). + # Only module-level built-ins get cache=True (see _build). + BinaryOp.register_new("floordiv", _floordiv, lazy=True, _cache=True) # cast to integer + BinaryOp.register_new( + "rfloordiv", _rfloordiv, lazy=True, _cache=True + ) # cast to integer # For aggregators - BinaryOp.register_new("absfirst", _absfirst, lazy=True) - BinaryOp.register_new("abssecond", _abssecond, lazy=True) - BinaryOp.register_new("rpow", _rpow, lazy=True) + BinaryOp.register_new("absfirst", _absfirst, lazy=True, _cache=True) + BinaryOp.register_new("abssecond", _abssecond, lazy=True, _cache=True) + BinaryOp.register_new("rpow", _rpow, lazy=True, _cache=True) # For algorithms binary._delayed["binom"] = (_register_binom, {}) # Lazy with custom creation @@ -981,7 +1142,7 @@ def _initialize(cls): (binary.any, _second), ]: binop.orig_func = func - binop._numba_func = numba.njit(func) if _has_numba else None + binop._numba_func = numba.njit(func, error_model="numpy") if _has_numba else None binop._udt_types = {} binop._udt_ops = {} binary.any._numba_func = binary.second._numba_func @@ -995,15 +1156,12 @@ def _initialize(cls): binary.eq._udt_ops = {} binary.ne._udt_types = {} binary.ne._udt_ops = {} - # Element-wise arithmetic ops are auto-generated from per-field / - # per-element scalar ops. - if _has_numba: - for op_name in _BUILTIN_UDT_BINARY_OPS: - binop = getattr(binary, op_name, None) - if binop is not None: - binop._udt_types = {} - binop._udt_ops = {} - binop._custom_dtype = _udt_dtype + # Element-wise arithmetic ops on UDTs are auto-generated from + # per-field / per-element scalar ops; the attributes that enable this + # are seeded in ``__init__`` (keyed on ``_BUILTIN_UDT_BINARY_OPS``). + # Do not seed them here with ``getattr(binary, op_name)``: that would + # force lazily-registered UDF ops like ``floordiv`` to numba-compile + # for every dtype at import time (about half a second of startup here). cls._initialized = True def __init__( @@ -1015,8 +1173,10 @@ def __init__( is_positional=False, is_udt=False, numba_func=None, + ret_dtype=None, ): super().__init__(name, anonymous=anonymous) + self._ret_dtype = ret_dtype self._monoid = None self._commutes_to = None self._semiring_commutes_to = None @@ -1025,9 +1185,29 @@ def __init__( self._is_udt = is_udt self.is_positional = is_positional self._custom_dtype = None + # Set True in ``_build`` for the built-in UDFs whose ``.types`` was + # populated by typing-only inference; ``_build_deferred`` then compiles + # each dtype's typed op on first request. + self._defer_builds = False if is_udt: self._udt_types = {} # {(dtype, dtype): DataType} self._udt_ops = {} # {(dtype, dtype): TypedUserBinaryOp} + elif _has_numba and name in _BUILTIN_UDT_BINARY_OPS: + # Built-in arithmetic ops auto-lift to UDTs field-by-field. + # Seeded here rather than in ``_initialize`` so that delayed ops + # (e.g. ``floordiv``) are not force-built at import time. + self._udt_types = {} + self._udt_ops = {} + self._custom_dtype = _udt_dtype + + def _build_deferred(self, type_): + # For built-ins whose ``.types`` was populated by typing-only inference, + # compile and register the typed op for ``type_`` the first time it is + # requested. ``ret_type`` is read back from ``.types`` (never re-derived + # per dtype: the UINT64/BOOL heuristic needs the full sample pass). + if not self._defer_builds or type_ not in self.types: + return None + return _finalize_typed_binaryop(self, self._numba_func, self.name, type_, self.types[type_]) __call__ = TypedBuiltinBinaryOp.__call__ is_commutative = TypedBuiltinBinaryOp.is_commutative diff --git a/graphblas/core/operator/indexbinary.py b/graphblas/core/operator/indexbinary.py index a19f5ffec..8c0fb6f26 100644 --- a/graphblas/core/operator/indexbinary.py +++ b/graphblas/core/operator/indexbinary.py @@ -6,7 +6,7 @@ from ...exceptions import UdfParseError, check_status_carg from .. import _has_numba, ffi, lib from ..dtypes import _sample_values -from .base import OpBase, ParameterizedUdf, TypedOpBase +from .base import OpBase, ParameterizedUdf, TypedOpBase, _validate_ret_dtype _has_idxbinop = hasattr(lib, "GxB_IndexBinaryOp_new") @@ -17,7 +17,7 @@ _bool_to_int8, _compile_udf_for_udt, _get_udt_wrapper_indexbinary, - _resolve_udt_return_type, + _udt_ret_type, ) ffi_new = ffi.new @@ -221,7 +221,7 @@ class IndexBinaryOp(OpBase): no built-ins; all IndexBinaryOps are user-defined. """ - __slots__ = "orig_func", "_is_udt", "_numba_func" + __slots__ = "orig_func", "_is_udt", "_numba_func", "_ret_dtype" _module = indexbinary _modname = "indexbinary" _custom_dtype = None @@ -235,7 +235,7 @@ class IndexBinaryOp(OpBase): } @classmethod - def _build(cls, name, func, *, is_udt=False, anonymous=False): + def _build(cls, name, func, *, is_udt=False, anonymous=False, ret_dtype=None): if not _has_idxbinop: raise RuntimeError( "IndexBinaryOp requires SuiteSparse:GraphBLAS 9.4+ " @@ -245,10 +245,20 @@ def _build(cls, name, func, *, is_udt=False, anonymous=False): raise TypeError(f"UDF argument must be a function, not {type(func)}") if name is None: name = getattr(func, "__name__", "") + ret_dtype = _validate_ret_dtype( + ret_dtype, "indexbinary", is_udt=is_udt, parameterized=False + ) success = False - indexbinary_udf = numba.njit(func) + # Set on the Dispatcher, not just the cfunc wrapper; see the note in + # ``BinaryOp._build``. + indexbinary_udf = numba.njit(func, error_model="numpy") new_type_obj = cls( - name, func, anonymous=anonymous, is_udt=is_udt, numba_func=indexbinary_udf + name, + func, + anonymous=anonymous, + is_udt=is_udt, + numba_func=indexbinary_udf, + ret_dtype=ret_dtype, ) return_types = {} nt = numba.types @@ -333,7 +343,9 @@ def indexbinary_wrapper( ): # pragma: no cover (numba) z[0] = indexbinary_udf(x[0], ix, jx, y[0], iy, jy, theta[0]) - indexbinary_wrapper = numba.cfunc(wrapper_sig, nopython=True)(indexbinary_wrapper) + indexbinary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")( + indexbinary_wrapper + ) new_idxbinop = ffi_new("GxB_IndexBinaryOp*") check_status_carg( lib.GxB_IndexBinaryOp_new( @@ -387,12 +399,14 @@ def _compile_udt(self, dtype, dtype2): numba_func, sig, op_kind="indexbinary", op_name=self.name, dtypes=(dtype, dtype2) ) numba_ret_type = numba_func.overloads[sig].signature.return_type - ret_type = _resolve_udt_return_type(numba_ret_type, dtype, dtype2) + ret_type = _udt_ret_type(self, numba_ret_type, dtype, dtype2) indexbinary_wrapper, wrapper_sig = _get_udt_wrapper_indexbinary( numba_func, ret_type, dtype, dtype2, numba_ret_type=numba_ret_type ) - indexbinary_wrapper = numba.cfunc(wrapper_sig, nopython=True)(indexbinary_wrapper) + indexbinary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")( + indexbinary_wrapper + ) new_idxbinop = ffi_new("GxB_IndexBinaryOp*") check_status_carg( lib.GxB_IndexBinaryOp_new( @@ -421,7 +435,9 @@ def _compile_udt(self, dtype, dtype2): return op @classmethod - def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=False): + def register_anonymous( + cls, func, name=None, *, parameterized=False, is_udt=False, ret_dtype=None + ): """Register an IndexBinaryOp without adding it to the ``indexbinary`` namespace. Because it is not registered in the namespace, the name is optional. @@ -442,6 +458,13 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals is_udt : bool, default False Whether the operator is intended to operate on user-defined types. + ret_dtype : dtype, optional + The dtype the operator returns. Requires ``is_udt=True``. + Without it the return type is inferred from what the function + returns, which can only name a type that is already an input, so + an output UDT that is not an operand needs this. The dtype is + fixed for the operator: it is the same for every input dtype. + Returns ------- IndexBinaryOp or ParameterizedIndexBinaryOp @@ -449,11 +472,14 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals """ cls._check_supports_udf("register_anonymous") if parameterized: + _validate_ret_dtype(ret_dtype, "indexbinary", is_udt=is_udt, parameterized=True) return ParameterizedIndexBinaryOp(name, func, anonymous=True, is_udt=is_udt) - return cls._build(name, func, anonymous=True, is_udt=is_udt) + return cls._build(name, func, anonymous=True, is_udt=is_udt, ret_dtype=ret_dtype) @classmethod - def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=False): + def register_new( + cls, name, func, *, parameterized=False, is_udt=False, lazy=False, ret_dtype=None + ): """Register a new IndexBinaryOp under the ``graphblas.indexbinary`` namespace. Parameters @@ -473,23 +499,37 @@ def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=Fal lazy : bool, default False When True, defer compilation until the operator is first used. + ret_dtype : dtype, optional + The dtype the operator returns. Requires ``is_udt=True``. + See :meth:`register_anonymous` for details. + Examples -------- >>> gb.indexbinary.register_new("index_dist", lambda x, ix, jx, y, iy, jy, t: abs(ix - iy)) """ cls._check_supports_udf("register_new") + # Validate eagerly even for lazy=True, so a bad combination fails at + # the registration site rather than at first attribute touch. + _validate_ret_dtype(ret_dtype, "indexbinary", is_udt=is_udt, parameterized=parameterized) module, funcname = cls._remove_nesting(name) if lazy: module._delayed[funcname] = ( cls.register_new, - {"name": name, "func": func, "parameterized": parameterized, "is_udt": is_udt}, + { + "name": name, + "func": func, + "parameterized": parameterized, + "is_udt": is_udt, + "ret_dtype": ret_dtype, + }, ) elif parameterized: + _validate_ret_dtype(ret_dtype, "indexbinary", is_udt=is_udt, parameterized=True) idxbinop = ParameterizedIndexBinaryOp(name, func, is_udt=is_udt) setattr(module, funcname, idxbinop) else: - idxbinop = cls._build(name, func, is_udt=is_udt) + idxbinop = cls._build(name, func, is_udt=is_udt, ret_dtype=ret_dtype) setattr(module, funcname, idxbinop) if not cls._initialized: # pragma: no cover (safety) @@ -505,8 +545,11 @@ def _initialize(cls): # No built-in IndexBinaryOps to register. cls._initialized = True - def __init__(self, name, func=None, *, anonymous=False, is_udt=False, numba_func=None): + def __init__( + self, name, func=None, *, anonymous=False, is_udt=False, numba_func=None, ret_dtype=None + ): super().__init__(name, anonymous=anonymous) + self._ret_dtype = ret_dtype self.orig_func = func self._numba_func = numba_func self._is_udt = is_udt diff --git a/graphblas/core/operator/indexunary.py b/graphblas/core/operator/indexunary.py index c7af8ed78..d4d0da79d 100644 --- a/graphblas/core/operator/indexunary.py +++ b/graphblas/core/operator/indexunary.py @@ -7,7 +7,7 @@ from ...exceptions import UdfParseError, check_status_carg from .. import _has_numba, ffi, lib from ..dtypes import _sample_values -from .base import OpBase, ParameterizedUdf, TypedOpBase, _call_op +from .base import OpBase, ParameterizedUdf, TypedOpBase, _call_op, _validate_ret_dtype if _has_numba: import numba @@ -17,7 +17,7 @@ _compile_udf_for_udt, _finalize_udt_op, _get_udt_wrapper, - _resolve_udt_return_type, + _udt_ret_type, ) ffi_new = ffi.new @@ -85,7 +85,7 @@ class IndexUnaryOp(OpBase): Built-in and registered IndexUnaryOps are located in the ``graphblas.indexunary`` namespace. """ - __slots__ = "orig_func", "is_positional", "_is_udt", "_numba_func" + __slots__ = "orig_func", "is_positional", "_is_udt", "_numba_func", "_ret_dtype" _module = indexunary _modname = "indexunary" _custom_dtype = None @@ -110,15 +110,23 @@ class IndexUnaryOp(OpBase): "rowindex", "colindex"} # fmt: skip @classmethod - def _build(cls, name, func, *, is_udt=False, anonymous=False): + def _build(cls, name, func, *, is_udt=False, anonymous=False, ret_dtype=None): if not isinstance(func, FunctionType): raise TypeError(f"UDF argument must be a function, not {type(func)}") if name is None: name = getattr(func, "__name__", "") + ret_dtype = _validate_ret_dtype(ret_dtype, "indexunary", is_udt=is_udt, parameterized=False) success = False - indexunary_udf = numba.njit(func) + # Set on the Dispatcher, not just the cfunc wrapper; see the note in + # ``BinaryOp._build``. + indexunary_udf = numba.njit(func, error_model="numpy") new_type_obj = cls( - name, func, anonymous=anonymous, is_udt=is_udt, numba_func=indexunary_udf + name, + func, + anonymous=anonymous, + is_udt=is_udt, + numba_func=indexunary_udf, + ret_dtype=ret_dtype, ) return_types = {} nt = numba.types @@ -176,7 +184,9 @@ def indexunary_wrapper(z, x, row, col, y): # pragma: no cover (numba) def indexunary_wrapper(z, x, row, col, y): # pragma: no cover (numba) z[0] = indexunary_udf(x[0], row, col, y[0]) - indexunary_wrapper = numba.cfunc(wrapper_sig, nopython=True)(indexunary_wrapper) + indexunary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")( + indexunary_wrapper + ) new_indexunary = ffi_new("GrB_IndexUnaryOp*") check_status_carg( lib.GrB_IndexUnaryOp_new( @@ -212,7 +222,7 @@ def _compile_udt(self, dtype, dtype2): numba_func, sig, op_kind="indexunary", op_name=self.name, dtypes=(dtype, dtype2) ) numba_ret_type = numba_func.overloads[sig].signature.return_type - ret_type = _resolve_udt_return_type(numba_ret_type, dtype, dtype2) + ret_type = _udt_ret_type(self, numba_ret_type, dtype, dtype2) indexunary_wrapper, wrapper_sig = _get_udt_wrapper( numba_func, ret_type, dtype, dtype2, include_indexes=True, numba_ret_type=numba_ret_type ) @@ -221,7 +231,9 @@ def _compile_udt(self, dtype, dtype2): ) @classmethod - def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=False): + def register_anonymous( + cls, func, name=None, *, parameterized=False, is_udt=False, ret_dtype=None + ): """Register a IndexUnary without registering it in the ``graphblas.indexunary`` namespace. Because it is not registered in the namespace, the name is optional. @@ -255,6 +267,13 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals Setting ``is_udt=True`` is also helpful when the left and right dtypes need to be different. + ret_dtype : dtype, optional + The dtype the operator returns. Requires ``is_udt=True``. + Without it the return type is inferred from what the function + returns, which can only name a type that is already an input, so + an output UDT that is not an operand needs this. The dtype is + fixed for the operator: it is the same for every input dtype. + Returns ------- return IndexUnaryOp or ParameterizedIndexUnaryOp @@ -262,11 +281,14 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals """ cls._check_supports_udf("register_anonymous") if parameterized: + _validate_ret_dtype(ret_dtype, "indexunary", is_udt=is_udt, parameterized=True) return ParameterizedIndexUnaryOp(name, func, anonymous=True, is_udt=is_udt) - return cls._build(name, func, anonymous=True, is_udt=is_udt) + return cls._build(name, func, anonymous=True, is_udt=is_udt, ret_dtype=ret_dtype) @classmethod - def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=False): + def register_new( + cls, name, func, *, parameterized=False, is_udt=False, lazy=False, ret_dtype=None + ): """Register a new IndexUnaryOp and save it to ``graphblas.indexunary`` namespace. If the return type is Boolean, the function will also be registered as a SelectOp @@ -309,6 +331,10 @@ def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=Fal delay compilation and only compile when the operator is used, which is done by setting ``lazy=True``. + ret_dtype : dtype, optional + The dtype the operator returns. Requires ``is_udt=True``. + See :meth:`register_anonymous` for details. + Examples -------- >>> gb.indexunary.register_new("row_mod", lambda x, i, j, thunk: i % max(thunk, 2)) @@ -317,17 +343,27 @@ def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=Fal """ cls._check_supports_udf("register_new") + # Validate eagerly even for lazy=True, so a bad combination fails at + # the registration site rather than at first attribute touch. + _validate_ret_dtype(ret_dtype, "indexunary", is_udt=is_udt, parameterized=parameterized) module, funcname = cls._remove_nesting(name) if lazy: module._delayed[funcname] = ( cls.register_new, - {"name": name, "func": func, "parameterized": parameterized, "is_udt": is_udt}, + { + "name": name, + "func": func, + "parameterized": parameterized, + "is_udt": is_udt, + "ret_dtype": ret_dtype, + }, ) elif parameterized: + _validate_ret_dtype(ret_dtype, "indexunary", is_udt=is_udt, parameterized=True) indexunary_op = ParameterizedIndexUnaryOp(name, func, is_udt=is_udt) setattr(module, funcname, indexunary_op) else: - indexunary_op = cls._build(name, func, is_udt=is_udt) + indexunary_op = cls._build(name, func, is_udt=is_udt, ret_dtype=ret_dtype) setattr(module, funcname, indexunary_op) # If return type is BOOL, register additionally as a SelectOp if all(x == BOOL for x in indexunary_op.types.values()): @@ -391,8 +427,10 @@ def __init__( is_positional=False, is_udt=False, numba_func=None, + ret_dtype=None, ): super().__init__(name, anonymous=anonymous) + self._ret_dtype = ret_dtype self.orig_func = func self._numba_func = numba_func self.is_positional = is_positional diff --git a/graphblas/core/operator/select.py b/graphblas/core/operator/select.py index 32a2509a5..7164694e4 100644 --- a/graphblas/core/operator/select.py +++ b/graphblas/core/operator/select.py @@ -97,10 +97,13 @@ def _from_indexunary(cls, iop): t.return_type, t.gb_obj, ) - # Aliases the IndexUnaryOp's allocation. The IndexUnaryOp - # owns the free; clearing here prevents a double free when - # both ends are GC'd. - op._owns_gb_obj_inst = False + # Borrow the IndexUnaryOp's allocation instead of making a + # second one. Holding ``t`` keeps that handle alive for as long + # as this SelectOp can use it: ``iop`` is a temporary in + # ``register_anonymous``, so without this the handle is freed + # the moment it is collected and every call raises + # UninitializedObject. + op._gb_obj_owner = t else: op = cls._typed_class( obj, diff --git a/graphblas/core/operator/semiring.py b/graphblas/core/operator/semiring.py index 32e20d0bf..336768a76 100644 --- a/graphblas/core/operator/semiring.py +++ b/graphblas/core/operator/semiring.py @@ -244,6 +244,9 @@ def _build(cls, name, monoid, binaryop, *, anonymous=False): new_type_obj = cls(name, monoid, binaryop, anonymous=anonymous) if binaryop._is_udt: return new_type_obj + # A built-in UDF multiplier (e.g. floordiv) may defer its per-dtype + # builds; force them so the iteration below sees every typed op. + binaryop._materialize_deferred() for binary_in, binary_func in binaryop._typed_ops.items(): binary_out = binary_func.return_type # Unfortunately, we can't have user-defined monoids over bools yet diff --git a/graphblas/core/operator/udt_utils.py b/graphblas/core/operator/udt_utils.py index fc33ea8c0..6625483ce 100644 --- a/graphblas/core/operator/udt_utils.py +++ b/graphblas/core/operator/udt_utils.py @@ -68,7 +68,12 @@ def _compile_codegen(src, *, func_name, source_label, extra_ns=None): filename, ) code = compile(src, filename, "exec") - namespace = {"min": min, "max": max, "abs": abs} + # No ``min`` / ``max`` here on purpose. Binding Python's builtins is what + # gave UDT ``binary.min`` its order-dependent NaN handling, and no name in + # this namespace has the C99 ``fmin`` semantics SuiteSparse uses, so + # :func:`_minmax_expr` spells the comparison out instead. ``signbit`` is + # the one piece it needs that isn't syntax. + namespace = {"abs": abs, "signbit": np.signbit} if _has_numba: namespace["numba"] = numba if extra_ns: @@ -123,8 +128,39 @@ def _compile_codegen(src, *, func_name, source_label, extra_ns=None): # (no ordering for min/max, no integer-mod for floordiv). _OPS_NOT_FOR_COMPLEX = frozenset({"min", "max", "floordiv"}) -# C operator equivalents for JIT code generation -_C_INFIX_OPS = {"+": "+", "-": "-", "*": "*", "/": "/", "//": "/"} +# ``_MIN`` spellings from , keyed by itemsize. Signed integer +# division traps on ``MIN / -1``; the generated C tests for it explicitly. +_C_INT_MIN = {1: "INT8_MIN", 2: "INT16_MIN", 4: "INT32_MIN", 8: "INT64_MIN"} + +# Same constants for the generated Python source, spelled ``-MAX - 1`` rather +# than as a direct literal that reads as the negation of an out-of-range value. +# MAINT 2026-07-30: Numba 0.65 types both spellings as ``Literal[int]``, so this +# can likely be simplified; re-check on the oldest supported Numba first. +_PY_INT_MIN = { + 1: "(-127 - 1)", + 2: "(-32767 - 1)", + 4: "(-2147483647 - 1)", + 8: "(-9223372036854775807 - 1)", +} + + +def _is_float(np_dtype): + return np_dtype.kind == "f" + + +def _is_signed_int(np_dtype): + return np_dtype.kind == "i" + + +def _is_int(np_dtype): + """True for signed and unsigned integers, but not bool. + + Bool is excluded from the integer division guards: ``_Bool`` has no range + to overflow and C converts a double to it by comparing against zero, so + both execution paths already agree without help. + """ + return np_dtype.kind in ("i", "u") + # Vanilla strips GxB callables but keeps GxB constants, so the bare # ``hasattr`` would lie; gate on the backend too. @@ -325,12 +361,27 @@ def _check_udt_pair(op_name, dtype, dtype2, info_x, info_y): f"binary.{op_name} does not work with ({dtype}, {dtype2}): " f"cannot mix record and array UDTs in a single element-wise op." ) - if kind_x == "record" and detail_x != detail_y: - raise KeyError( - f"binary.{op_name} does not work with ({dtype}, {dtype2}): " - f"record UDTs must share field names; got {list(detail_x)} vs " - f"{list(detail_y)}." - ) + if kind_x == "record": + if detail_x != detail_y: + raise KeyError( + f"binary.{op_name} does not work with ({dtype}, {dtype2}): " + f"record UDTs must share field names; got {list(detail_x)} vs " + f"{list(detail_y)}." + ) + # Matching top-level names is not enough: the codegen pairs operands + # leaf by leaf, and a field that is a sub-record on one side and a + # scalar on the other contributes a different number of leaves. Left + # unchecked the pair reaches Numba, whose typing failure arrives as a + # UdfParseError, reporting a compile error for what is really the same + # shape disagreement the checks above report as a KeyError. + leaves_x = [c for _py, c, _d in _iter_record_leaves(dtype.np_type)] + leaves_y = [c for _py, c, _d in _iter_record_leaves(dtype2.np_type)] + if len(leaves_x) != len(leaves_y): + raise KeyError( + f"binary.{op_name} does not work with ({dtype}, {dtype2}): " + f"record UDTs must nest the same way, so that each has the same " + f"number of leaf fields; got {leaves_x} vs {leaves_y}." + ) if kind_x == "array" and detail_x != detail_y: raise KeyError( f"binary.{op_name} does not work with ({dtype}, {dtype2}): " @@ -349,9 +400,11 @@ def _iter_record_leaves(np_type, python_prefix="", c_prefix=""): ``("['outer']['inner_a']", "outer.inner_a", float64_dtype)``. Only record-in-record nesting is recognised. Array-typed fields are - yielded as a single leaf (the codegen treats them opaquely; arithmetic - on an array-typed sub-field isn't supported by either the cfunc or the - JIT path today). + yielded as a single leaf, which the Numba path handles: the generated + expression operates on the whole sub-array and the wrapper slice-assigns + it back. The JIT path never sees such a record, because a subarray dtype + has no entry in ``NP_TO_C_TYPES`` and ``_udt_c_typedef`` bails out, so + the op keeps the cfunc for every call. """ for name in np_type.names: field_dtype = np_type.fields[name][0] @@ -513,85 +566,209 @@ def _op_supports_field_dtypes(op_name, np_type): # Numba function generators, called lazily from each op's ``_compile_udt``. if _has_numba: - def _expr_binary(py_op, x_expr, y_expr): - """Python-source builder; sibling of :func:`_c_expr_binary` for JIT C.""" + def _expr_binary(py_op, x_expr, y_expr, x_dtype, y_dtype): + """Python-source builder; sibling of :func:`_c_expr_binary` for JIT C. + + ``x_dtype`` and ``y_dtype`` are each operand's numpy dtype at this + leaf, which ``_check_udt_pair`` allows to differ. Four ops consult + them: + + - ``min`` / ``max`` follow SuiteSparse's own ``GrB_MIN_FP64``, which + is C99 ``fmin`` and ignores a NaN operand rather than ordering or + propagating it; see :func:`_minmax_expr`. Only floating-point + leaves carry the NaN and signed-zero rules, so the dtype decides + which of the two forms is emitted. + - ``//`` on two signed integers. Numba deliberately returns 0 for + ``INT_MIN // -1`` to dodge the SIGFPE that x86 raises on the + unrepresentable quotient, while numpy wraps to ``INT_MIN``. Since + ``a // -1`` is exactly ``-a``, routing that divisor through + negation (which wraps) reaches numpy's answer without the trap. + - ``/`` on two integers. The quotient is a float that then has to fit + back into the integer field; when it doesn't, the conversion is + undefined in both C and LLVM, and the two disagree in ways that + vary with the field's width. Ruling out the two ways a division of + two same-signedness integers can leave the field's range (a zero + divisor, and ``INT_MIN / -1``) is what makes the two paths agree. + Note this gives ``x / 0 == 0``, which agrees with + ``np.floor_divide``; ``np.true_divide`` produces an infinity whose + cast back to an integer is undefined in numpy too. + + Two range escapes are left, and both move the two paths together + rather than apart. A 64-bit field can leave the range through the + ``(double)`` conversion itself (``(2**63 - 1) / 1`` rounds up to + ``2**63``), and there the answer is whatever the hardware does: + measured saturating to ``INT64_MAX`` on arm64. Operands of mixed + signedness escape through a negative divisor, which the + ``INT_MIN`` guard, keyed on the left operand, doesn't see: + ``uint8(200) / int8(-1)`` is -200.0 before the cast back. That one + is unguarded rather than known-wrong (it lands on 56, which is what + ``np.float64(-200).astype(np.uint8)`` gives), and a mixed pair gets + no JIT kernel at all, so there is no second path for it to + disagree with. + - ``/`` where either side is complex. Numba's complex division raises + ``ZeroDivisionError`` unconditionally, outside the error model's + control, so a zero divisor left the element unwritten while the C + kernel returned numpy's infinities. Spelling out the zero case here + makes both paths match numpy. + """ if py_op in _FUNC_BINARY_OPS: - return f"{py_op}({x_expr}, {y_expr})" + return _minmax_expr(py_op, x_expr, y_expr, x_dtype, y_dtype) + if py_op == "//" and _is_signed_int(x_dtype) and _is_signed_int(y_dtype): + return f"(-{x_expr} if {y_expr} == -1 else {x_expr} // {y_expr})" + if py_op == "/" and (x_dtype.kind == "c" or y_dtype.kind == "c"): + return ( + f"({x_expr} / {y_expr} if {y_expr} != 0 " + f"else complex({x_expr}.real / 0.0, {x_expr}.imag / 0.0))" + ) + if py_op == "/" and _is_int(x_dtype) and _is_int(y_dtype): + quotient = f"{x_expr} / {y_expr}" + if _is_signed_int(x_dtype): + type_min = _PY_INT_MIN[x_dtype.itemsize] + quotient = ( + f"({x_expr} if ({x_expr} == {type_min} and {y_expr} == -1) else {quotient})" + ) + return f"(0 if {y_expr} == 0 else {quotient})" return f"{x_expr} {py_op} {y_expr}" + def _minmax_expr(py_op, x_expr, y_expr, x_dtype, y_dtype): + """Return a Python expression for ``min`` / ``max`` matching C ``fmin``. + + ``GrB_MIN_FP64`` is C99 ``fmin``, so this builder and + :func:`_c_minmax_expr` both have to reproduce it or ``binary.min`` + means one thing on FP64 and another on a UDT. Two rules follow from + that: a NaN operand is ignored unless both are NaN, and a tie between + ``-0.0`` and ``0.0`` resolves to ``-0.0`` for ``min``, ``0.0`` for + ``max``, whichever side each sits on. + + Neither ``min`` nor ``np.fmin`` gets both rules right under Numba. + Python's builtin returns its first argument whenever the comparison + is false, so it keeps a NaN on the left and drops one on the right. + ``np.fmin`` fixes the NaN rule but Numba lowers it without the + signed-zero tie-break that libm has, so it would disagree with the + JIT C kernel on the sign of a zero. Spelling the comparison out is + what keeps the two execution paths equal. + + Integers have neither NaN nor a signed zero, so they take the same + plain comparison the C side emits. A leaf is treated as + floating-point if *either* operand is: broadcasting a float scalar + over an int record still produces float results. + """ + cmp_op = "<" if py_op == "min" else ">" + if not _is_float(x_dtype) and not _is_float(y_dtype): + return f"({x_expr} if {x_expr} {cmp_op} {y_expr} else {y_expr})" + # ``y != y`` is the NaN test on the right operand: when it holds, the + # left one wins whatever it is, which is fmin's "ignore the NaN" rule + # and also (correctly) returns NaN when both are NaN. + tie = f"signbit({x_expr})" if py_op == "min" else f"not signbit({x_expr})" + return ( + f"({x_expr} if ({x_expr} {cmp_op} {y_expr} or {y_expr} != {y_expr}" + f" or ({x_expr} == {y_expr} and {tie})) else {y_expr})" + ) + def _expr_unary(py_op, operand): """Python-source builder; sibling of :func:`_c_expr_unary` for JIT C.""" if py_op in _FUNC_UNARY_OPS: return f"{py_op}({operand})" return f"{py_op}{operand}" - def _make_record_func(leaf_paths, arity, py_op, *, x_is_scalar=False, y_is_scalar=False): + def _make_record_func( + leaves, + arity, + py_op, + *, + x_is_scalar=False, + y_is_scalar=False, + x_scalar_dtype=None, + y_scalar_dtype=None, + ): """Build a Numba njit function for a record UDT. - ``leaf_paths`` is a sequence of Python access strings ``"['a']"`` or, - for nested records, ``"['outer']['inner_a']"``. The generated function - always returns a *flat* tuple of leaf values regardless of nesting - depth; the wrapper (in base.py) walks the same leaf paths when - writing the result back, so nested-record outputs land at the - correct depth without nested tuple construction (which Numba can't - ``setitem``-assign to a record field). + ``leaves`` is a sequence of ``(python_access, x_leaf_dtype, + y_leaf_dtype)`` triples, where ``y_leaf_dtype`` is ``None`` for a unary + op. The two dtypes differ when the operands are records that share + field names but not field types, which ``_check_udt_pair`` allows. + The access strings look like ``"['a']"`` or, for nested records, + ``"['outer']['inner_a']"``. The generated function always returns a + *flat* tuple of leaf values regardless of nesting depth; the wrapper + (in base.py) walks the same leaf paths when writing the result back, + so nested-record outputs land at the correct depth without nested + tuple construction (which Numba can't ``setitem``-assign to a record + field). When ``x_is_scalar`` or ``y_is_scalar`` is True, that argument is a - plain scalar (not a record), so it is used directly for all leaves. + plain scalar (not a record), so it is used directly for all leaves + and ``x_scalar_dtype`` / ``y_scalar_dtype`` gives its numpy dtype. """ if arity == 2: parts = [] - for path in leaf_paths: + for path, x_leaf_dtype, y_leaf_dtype in leaves: x_expr = "x" if x_is_scalar else f"x{path}" y_expr = "y" if y_is_scalar else f"y{path}" - parts.append(_expr_binary(py_op, x_expr, y_expr)) + parts.append( + _expr_binary( + py_op, + x_expr, + y_expr, + x_scalar_dtype if x_is_scalar else x_leaf_dtype, + y_scalar_dtype if y_is_scalar else y_leaf_dtype, + ) + ) sig = "x, y" else: - parts = [_expr_unary(py_op, f"x{path}") for path in leaf_paths] + parts = [_expr_unary(py_op, f"x{path}") for path, _xd, _yd in leaves] sig = "x" body = ", ".join(parts) # Single-leaf tuple needs the trailing comma to remain a tuple. - ret = f"({body},)" if len(leaf_paths) == 1 else f"({body})" + ret = f"({body},)" if len(leaves) == 1 else f"({body})" src = f"def _op({sig}):\n return {ret}\n" op_func = _compile_codegen( src, func_name="_op", - source_label=f"", + source_label=f"", ) - return numba.njit(op_func) + return numba.njit(op_func, error_model="numpy") def _make_array_wrapper( size, - base_numba_type, + base_dtype, arity, py_op, *, - x_scalar_type=None, - y_scalar_type=None, + x_scalar_dtype=None, + y_scalar_dtype=None, ): """Build a cfunc-ready wrapper for an array UDT (element-by-element). - When ``x_scalar_type`` or ``y_scalar_type`` is set, that side is a plain - scalar pointer (broadcast to all elements). + All dtype arguments are numpy dtypes. When ``x_scalar_dtype`` or + ``y_scalar_dtype`` is set, that side is a plain scalar pointer + (broadcast to all elements). Returns (wrapper_func, wrapper_sig). """ nt = numba.types + base_numba_type = numba.from_dtype(base_dtype) if arity == 2: - x_ref = "x_ptr[0]" if x_scalar_type else "x[{i}]" - y_ref = "y_ptr[0]" if y_scalar_type else "y[{i}]" - assigns = "\n".join( - f" z[{i}] = {_expr_binary(py_op, x_ref.format(i=i), y_ref.format(i=i))}" - for i in range(size) - ) + x_ref = "x_ptr[0]" if x_scalar_dtype is not None else "x[{i}]" + y_ref = "y_ptr[0]" if y_scalar_dtype is not None else "y[{i}]" + x_dtype = base_dtype if x_scalar_dtype is None else x_scalar_dtype + y_dtype = base_dtype if y_scalar_dtype is None else y_scalar_dtype + lines = [] + for i in range(size): + expr = _expr_binary(py_op, x_ref.format(i=i), y_ref.format(i=i), x_dtype, y_dtype) + lines.append(f" z[{i}] = {expr}") + assigns = "\n".join(lines) params = "z_ptr, x_ptr, y_ptr" arrays = f" z = numba.carray(z_ptr, {size})\n" - if not x_scalar_type: + if x_scalar_dtype is None: arrays += f" x = numba.carray(x_ptr, {size})\n" - if not y_scalar_type: + if y_scalar_dtype is None: arrays += f" y = numba.carray(y_ptr, {size})\n" - x_numba = nt.CPointer(x_scalar_type) if x_scalar_type else nt.CPointer(base_numba_type) - y_numba = nt.CPointer(y_scalar_type) if y_scalar_type else nt.CPointer(base_numba_type) + x_numba = nt.CPointer( + base_numba_type if x_scalar_dtype is None else numba.from_dtype(x_scalar_dtype) + ) + y_numba = nt.CPointer( + base_numba_type if y_scalar_dtype is None else numba.from_dtype(y_scalar_dtype) + ) sig = nt.void(nt.CPointer(base_numba_type), x_numba, y_numba) else: assigns = "\n".join( @@ -672,20 +849,40 @@ def compile_udt_binary_wrapper(op_name, py_op, dtype, dtype2): # Use leaf paths so the same codegen handles nested-record UDTs # uniformly. A non-nested record's leaves are its top-level # fields, with paths like ``"['a']"``. - leaf_paths = [py for py, _c, _d in _iter_record_leaves(udt_dtype.np_type)] + # + # Pair each leaf with its own operand's dtype. ``_check_udt_pair`` + # makes two record operands share field names but not field types, + # so reusing the left record's dtypes for both would describe the + # right operand's leaves incorrectly. + x_leaves = _iter_record_leaves((dtype if not x_is_scalar else udt_dtype).np_type) + y_leaves = _iter_record_leaves((dtype2 if not y_is_scalar else udt_dtype).np_type) + leaves = [ + (py, x_leaf_dtype, y_leaf_dtype) + for (py, _cx, x_leaf_dtype), (_py, _cy, y_leaf_dtype) in zip( + x_leaves, y_leaves, strict=True + ) + ] func = _make_record_func( - leaf_paths, + leaves, 2, py_op, x_is_scalar=x_is_scalar, y_is_scalar=y_is_scalar, + x_scalar_dtype=dtype.np_type if x_is_scalar else None, + y_scalar_dtype=dtype2.np_type if y_is_scalar else None, ) sig = (dtype.numba_type, dtype2.numba_type) _compile_udf_for_udt( func, sig, op_kind="binary", op_name=op_name, dtypes=(dtype, dtype2) ) numba_ret_type = func.overloads[sig].signature.return_type - ret_type = _resolve_udt_return_type(numba_ret_type, udt_dtype) + # Offer both operands when both are UDTs: passing only ``udt_dtype`` + # left the resolver no choice but the left-hand record, so an + # int-record combined with a float-record truncated to the int one + # (and gave a different answer if you swapped the operands). + ret_type = _resolve_udt_return_type( + numba_ret_type, *(d for d in (dtype, dtype2) if d._is_udt) + ) wrapper, wrapper_sig = _get_udt_wrapper( func, ret_type, dtype, dtype2, numba_ret_type=numba_ret_type ) @@ -694,11 +891,11 @@ def compile_udt_binary_wrapper(op_name, py_op, dtype, dtype2): ret_type = udt_dtype wrapper, wrapper_sig = _make_array_wrapper( size, - numba.from_dtype(base_dtype), + base_dtype, 2, py_op, - x_scalar_type=numba.from_dtype(dtype.np_type) if x_is_scalar else None, - y_scalar_type=numba.from_dtype(dtype2.np_type) if y_is_scalar else None, + x_scalar_dtype=dtype.np_type if x_is_scalar else None, + y_scalar_dtype=dtype2.np_type if y_is_scalar else None, ) return wrapper, wrapper_sig, ret_type @@ -722,8 +919,8 @@ def compile_udt_unary_wrapper(op_name, py_op, dtype): if kind == "record": from .base import _compile_udf_for_udt - leaf_paths = [py for py, _c, _d in _iter_record_leaves(dtype.np_type)] - func = _make_record_func(leaf_paths, 1, py_op) + leaves = [(py, d, None) for py, _c, d in _iter_record_leaves(dtype.np_type)] + func = _make_record_func(leaves, 1, py_op) sig = (dtype.numba_type,) _compile_udf_for_udt(func, sig, op_kind="unary", op_name=op_name, dtypes=(dtype,)) numba_ret_type = func.overloads[sig].signature.return_type @@ -734,64 +931,194 @@ def compile_udt_unary_wrapper(op_name, py_op, dtype): else: base_dtype, size = detail ret_type = dtype - wrapper, wrapper_sig = _make_array_wrapper(size, numba.from_dtype(base_dtype), 1, py_op) + wrapper, wrapper_sig = _make_array_wrapper(size, base_dtype, 1, py_op) return wrapper, wrapper_sig, ret_type # JIT C code generators below. -def _c_expr_binary(py_op, lhs, rhs, field_dtype=None): - """Return a C expression for a binary op: e.g., ``(x->a) + (y->a)``. +def _c_assign_binary(py_op, target, lhs, rhs, field_dtype): + """Return a C *statement* assigning ``lhs py_op rhs`` to ``target``. - ``field_dtype`` is the numpy dtype of the *result* element. It is only - consulted for ``floordiv`` (``//``), which needs Python ``//`` semantics - rather than C ``/`` (trunc toward zero for ints, true division for - floats). Other ops are type-agnostic at the C level. + Almost every op is a single expression, but floating-point ``//`` needs + temporaries (see :func:`_c_float_floordiv_stmt`), so this is the entry + point the kernel builder uses rather than :func:`_c_expr_binary`. """ - if py_op == "min": - # Match Python ``min(a, b) = b if b < a else a`` so NaN propagates - # from the first operand (cfunc / numba follows the same rule). - # The naive ``(a < b ? a : b)`` would silently swallow NaN to the - # right-hand side and disagree with the cfunc path. - return f"(({rhs}) < ({lhs}) ? ({rhs}) : ({lhs}))" - if py_op == "max": - return f"(({rhs}) > ({lhs}) ? ({rhs}) : ({lhs}))" - if py_op == "//": - return _c_floordiv_expr(lhs, rhs, field_dtype) - c_op = _C_INFIX_OPS.get(py_op, py_op) - return f"({lhs}) {c_op} ({rhs})" + if py_op == "//" and field_dtype.kind == "f": + return _c_float_floordiv_stmt(target, lhs, rhs, field_dtype) + return f"{target} = {_c_expr_binary(py_op, lhs, rhs, field_dtype)} ;" -def _c_floordiv_expr(lhs, rhs, field_dtype): - """Return a C expression for Python-semantics floor division. +def _c_expr_binary(py_op, lhs, rhs, field_dtype): + """Return a C expression for a binary op: e.g., ``(x->a) + (y->a)``. - Python ``//`` is floor (rounds toward negative infinity); C ``/`` is - trunc toward zero for ints and true division for floats. The two only - agree for non-negative integer operands; for everything else the JIT - path silently disagreed with the Numba cfunc path before this helper. + ``field_dtype`` is the numpy dtype of the element being written. Four of + the ops need it because the C spelling that looks obvious disagrees with + what SuiteSparse's own operators compute: ``/`` is integer division in C + but true division in Python, ``//`` is neither, and ``min`` / ``max`` have + to ignore NaN the way ``fmin`` does. It is required rather than optional + so a new caller can't silently get the wrong kernel. - Float fields use ``floor()`` / ``floorf()`` from ````, which is - available in the JIT kernel via SuiteSparse's include chain - (``GraphBLAS.h`` -> ````). Signed integer fields use the - standard trunc-to-floor adjustment. Unsigned integers don't need - adjusting because both operands are non-negative. + Floating-point ``//`` has no single-expression form; go through + :func:`_c_assign_binary`. + """ + if py_op in _FUNC_BINARY_OPS: + return _c_minmax_expr(py_op, lhs, rhs, field_dtype) + if py_op == "/": + return _c_truediv_expr(lhs, rhs, field_dtype) + if py_op == "//": + return _c_int_floordiv_expr(lhs, rhs, field_dtype) + # Everything left (``+``, ``-``, ``*``) spells the same in C as in Python + # and needs no dtype-specific handling. + return f"({lhs}) {py_op} ({rhs})" + + +def _c_minmax_expr(py_op, lhs, rhs, field_dtype): + """Return a C expression for ``min`` / ``max`` matching SuiteSparse. + + ``GrB_MIN_FP64`` is C99 ``fmin``: it ignores a NaN operand rather than + ordering it. A UDT field has to do the same or ``binary.min`` means two + different things depending on the dtype it is typed for. Calling ``fmin`` + itself is the way to be sure of that, tie-break on signed zeros included. + ```` reaches the JIT kernel through ``GraphBLAS.h``, the same + route ``_c_float_floordiv_stmt`` relies on for ``floor`` and ``fmod``. + + ``float`` fields need ``fminf``: passing them to ``fmin`` would compute + in double and round on the way back, which costs nothing in accuracy for + a min but does cost a conversion per element. + + Integers have neither NaN nor a signed zero, so a plain comparison is + exact for them. ``_minmax_expr`` emits the matching Python source; the + two must agree, because SuiteSparse picks between the JIT kernel and the + Numba cfunc on its own. + """ + if _is_float(field_dtype): + fn = "fmin" if py_op == "min" else "fmax" + suffix = "f" if field_dtype.itemsize == 4 else "" + return f"{fn}{suffix}(({lhs}), ({rhs}))" + cmp_op = "<" if py_op == "min" else ">" + return f"(({lhs}) {cmp_op} ({rhs}) ? ({lhs}) : ({rhs}))" + + +def _c_truediv_expr(lhs, rhs, field_dtype): + """Return a C expression for Python-semantics true division. + + C ``/`` on two integers is integer division; Python's ``/`` always + divides in floating point and only then does the result land back in the + integer field. The difference is visible whenever the exact quotient + doesn't fit a double (``10**18 / 3`` is 333333333333333333 in C but + 333333333333333312 through float64) and it is what makes integer + division by zero trap: ``(double) 7 / 0`` is ``inf``, but ``7 / 0`` in + integers raises SIGFPE and takes the whole process down. + + The quotient then has to come back down into the field, and converting a + double that doesn't fit the destination is undefined in C and poison in + LLVM, so the two paths need not agree. The two range escapes that a + same-signedness division can take are ruled out instead, matching the + guards :func:`_expr_binary` emits for the cfunc. A 64-bit field can still + escape through the ``(double)`` conversion (``(2**63 - 1) / 1`` rounds up + to ``2**63``). Both paths land on the same hardware conversion there + rather than on defined behaviour, so they agree with each other but need + not agree across machines; measured saturating to ``INT64_MAX`` on arm64. + + A complex field spells the zero-divisor case out rather than leaning on + C99 Annex G. SuiteSparse's JIT compiles with ``-fcx-limited-range`` under + GCC, which replaces the Annex G division with the naive formula, so + ``z / 0`` came out ``nan+nanj`` on Linux while clang (no such flag) and + the cfunc gave numpy's infinities. Dividing the parts by ``0.0`` as reals + sidesteps the flag entirely and matches the cfunc's spelling (see + ``_expr_binary``). ``CMPLX`` rather than arithmetic on ``I``: under the + naive formula ``inf * I`` multiplies out to ``nan``, which is the exact + failure being avoided. ``rhs == 0`` on a ``_Complex`` operand compares + both parts, mirroring the cfunc's ``y != 0``. """ - if field_dtype is None: - # Caller didn't pass dtype info. The C ``/`` semantics match Python - # ``//`` for non-negative integer operands only. - return f"({lhs}) / ({rhs})" kind = field_dtype.kind + if kind == "c": + if field_dtype.itemsize == 8: + cmplx, creal, cimag, zero = "CMPLXF", "crealf", "cimagf", "0.0f" + else: + cmplx, creal, cimag, zero = "CMPLX", "creal", "cimag", "0.0" + return ( + f"(({rhs}) == 0 " + f"? {cmplx}({creal}({lhs}) / {zero}, {cimag}({lhs}) / {zero}) " + f": ({lhs}) / ({rhs}))" + ) if kind == "f": - if field_dtype.itemsize == 4: - return f"floorf((float)({lhs}) / (float)({rhs}))" - return f"floor((double)({lhs}) / (double)({rhs}))" - if kind in ("u", "b"): return f"({lhs}) / ({rhs})" - # Signed integer: trunc-toward-zero is one greater than floor when the - # signs of ``a`` and ``b`` differ and the division has a non-zero - # remainder; subtract 1 in that case. - return f"(({lhs}) / ({rhs}) - ((({lhs}) % ({rhs}) != 0) && ((({lhs}) < 0) != (({rhs}) < 0))))" + quotient = f"(double)({lhs}) / (double)({rhs})" + if kind == "b": + # C converts a double to _Bool by comparing against zero, so there is + # no range to leave and nothing to guard. + return f"({quotient})" + if kind == "u": + return f"(({rhs}) == 0 ? 0 : ({quotient}))" + type_min = _C_INT_MIN[field_dtype.itemsize] + return ( + f"(({rhs}) == 0 ? 0 : " + f"((({lhs}) == {type_min} && ({rhs}) == -1) ? {type_min} : ({quotient})))" + ) + + +def _c_int_floordiv_expr(lhs, rhs, field_dtype): + """Return a C expression for integer floor division. + + Three ways C ``/`` differs from what the cfunc computes: + + - It truncates toward zero; Python ``//`` floors. Subtract one when the + operands have different signs and the remainder is non-zero. + - It raises SIGFPE (process death, not an exception) when the divisor is + zero. Numba's numpy error model returns 0, as does ``np.floor_divide``. + - It raises SIGFPE on ``INT_MIN / -1``, whose true quotient is not + representable. numpy wraps to ``INT_MIN``, which is what the cfunc + reaches through negation (see ``_expr_binary``). + """ + if field_dtype.kind in ("u", "b"): + # Both operands are non-negative, so truncation already floors. + return f"(({rhs}) == 0 ? 0 : (({lhs}) / ({rhs})))" + type_min = _C_INT_MIN[field_dtype.itemsize] + floored = f"({lhs}) / ({rhs}) - ((({lhs}) % ({rhs}) != 0) && ((({lhs}) < 0) != (({rhs}) < 0)))" + return ( + f"(({rhs}) == 0 ? 0 : " + f"((({lhs}) == {type_min} && ({rhs}) == -1) ? {type_min} : ({floored})))" + ) + + +def _c_float_floordiv_stmt(target, lhs, rhs, field_dtype): + """Return a C block computing floating-point ``//`` into ``target``. + + ``floor(a / b)`` is not floor division. numpy and CPython both compute + the quotient from the remainder (``(a - fmod(a, b)) / b``) and snap it, + which is more accurate and handles infinities differently: ``1.0 // 0.1`` + is 9.0 but ``floor(1.0 / 0.1)`` is 10.0, and ``inf // 2.0`` is NaN but + ``floor(inf / 2.0)`` is ``inf``. This mirrors numpy's ``npy_divmod`` + (which CPython's ``float_divmod`` matches), so the JIT kernel, the Numba + cfunc, and ``np.floor_divide`` all agree. + + The sign fix-up needs the remainder twice and the quotient twice, so + this emits a statement with temporaries rather than one expression. + """ + is_f32 = field_dtype.itemsize == 4 + ctype = "float" if is_f32 else "double" + sfx = "f" if is_f32 else "" + half = "0.5f" if is_f32 else "0.5" + one = "1.0f" if is_f32 else "1.0" + zero = "0.0f" if is_f32 else "0.0" + return ( + f"{{ {ctype} gb_a = ({lhs}) ; {ctype} gb_b = ({rhs}) ; {ctype} gb_q ; " + # A zero divisor is the one case numpy answers straight from the + # division: +-inf, or NaN for 0/0. + f"if (gb_b == {zero}) {{ gb_q = gb_a / gb_b ; }} " + f"else {{ {ctype} gb_m = fmod{sfx} (gb_a, gb_b) ; " + f"{ctype} gb_d = (gb_a - gb_m) / gb_b ; " + # NaN compares false both ways, so a NaN remainder skips the + # adjustment and carries through to the quotient. + f"if (gb_m != {zero} && ((gb_b < {zero}) != (gb_m < {zero}))) {{ gb_d -= {one} ; }} " + f"if (gb_d != {zero}) {{ gb_q = floor{sfx} (gb_d) ; " + f"if (gb_d - gb_q > {half}) {{ gb_q += {one} ; }} }} " + f"else {{ gb_q = copysign{sfx} ({zero}, gb_a / gb_b) ; }} }} " + f"{target} = gb_q ; }}" + ) def _c_expr_unary(py_op, operand, field_dtype=None): @@ -861,7 +1188,7 @@ def _make_jit_c_definition(op_name, py_op, dtype, arity): # Pass the leaf dtype to the binary expression builder so # type-sensitive ops (currently floordiv) can emit correct C. assigns = " ".join( - f"z->{c} = {_c_expr_binary(py_op, f'x->{c}', f'y->{c}', leaf_dtype)} ;" + _c_assign_binary(py_op, f"z->{c}", f"x->{c}", f"y->{c}", leaf_dtype) for _py, c, leaf_dtype in leaves ) else: @@ -874,7 +1201,7 @@ def _make_jit_c_definition(op_name, py_op, dtype, arity): size = reduce(mul, shape) if arity == 2: assigns = " ".join( - f"z->v[{i}] = {_c_expr_binary(py_op, f'x->v[{i}]', f'y->v[{i}]', base_dtype)} ;" + _c_assign_binary(py_op, f"z->v[{i}]", f"x->v[{i}]", f"y->v[{i}]", base_dtype) for i in range(size) ) else: diff --git a/graphblas/core/operator/unary.py b/graphblas/core/operator/unary.py index 22822921a..26359dc69 100644 --- a/graphblas/core/operator/unary.py +++ b/graphblas/core/operator/unary.py @@ -28,6 +28,7 @@ ParameterizedUdf, TypedOpBase, _hasop, + _validate_ret_dtype, ) if _supports_complex: @@ -40,7 +41,7 @@ _compile_udf_for_udt, _finalize_udt_op, _get_udt_wrapper, - _resolve_udt_return_type, + _udt_ret_type, ) ffi_new = ffi.new @@ -136,7 +137,7 @@ class UnaryOp(OpBase): as well as in the ``graphblas.ops`` combined namespace. """ - __slots__ = "orig_func", "is_positional", "_is_udt", "_numba_func" + __slots__ = "orig_func", "is_positional", "_is_udt", "_numba_func", "_ret_dtype" _custom_dtype = None _module = unary _modname = "unary" @@ -170,14 +171,24 @@ class UnaryOp(OpBase): _positional = {"positioni", "positioni1", "positionj", "positionj1"} @classmethod - def _build(cls, name, func, *, anonymous=False, is_udt=False): + def _build(cls, name, func, *, anonymous=False, is_udt=False, ret_dtype=None): if type(func) is not FunctionType: raise TypeError(f"UDF argument must be a function, not {type(func)}") if name is None: name = getattr(func, "__name__", "") + ret_dtype = _validate_ret_dtype(ret_dtype, "unary", is_udt=is_udt, parameterized=False) success = False - unary_udf = numba.njit(func) - new_type_obj = cls(name, func, anonymous=anonymous, is_udt=is_udt, numba_func=unary_udf) + # Set on the Dispatcher, not just the cfunc wrapper; see the note in + # ``BinaryOp._build``. + unary_udf = numba.njit(func, error_model="numpy") + new_type_obj = cls( + name, + func, + anonymous=anonymous, + is_udt=is_udt, + numba_func=unary_udf, + ret_dtype=ret_dtype, + ) return_types = {} nt = numba.types if not is_udt: @@ -231,7 +242,9 @@ def unary_wrapper(z, x): def unary_wrapper(z, x): z[0] = unary_udf(x[0]) # pragma: no cover (numba) - unary_wrapper = numba.cfunc(wrapper_sig, nopython=True)(unary_wrapper) + unary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")( + unary_wrapper + ) new_unary = ffi_new("GrB_UnaryOp*") check_status_carg( lib.GrB_UnaryOp_new( @@ -281,7 +294,7 @@ def _compile_udt(self, dtype, dtype2): sig = (dtype.numba_type,) _compile_udf_for_udt(numba_func, sig, op_kind="unary", op_name=self.name, dtypes=(dtype,)) numba_ret_type = numba_func.overloads[sig].signature.return_type - ret_type = _resolve_udt_return_type(numba_ret_type, dtype) + ret_type = _udt_ret_type(self, numba_ret_type, dtype) unary_wrapper, wrapper_sig = _get_udt_wrapper( numba_func, ret_type, dtype, numba_ret_type=numba_ret_type ) @@ -290,7 +303,9 @@ def _compile_udt(self, dtype, dtype2): ) @classmethod - def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=False): + def register_anonymous( + cls, func, name=None, *, parameterized=False, is_udt=False, ret_dtype=None + ): """Register a UnaryOp without registering it in the ``graphblas.unary`` namespace. Because it is not registered in the namespace, the name is optional. @@ -316,6 +331,12 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals Whether the operator is intended to operate on user-defined types. If True, then the function will not be automatically compiled for builtin types, and it will be compiled "just in time" when used. + ret_dtype : dtype, optional + The dtype the operator returns. Requires ``is_udt=True``. + Without it the return type is inferred from what the function + returns, which can only name a type that is already an input, so + an output UDT that is not an operand needs this. The dtype is + fixed for the operator: it is the same for every input dtype. Returns ------- @@ -324,11 +345,14 @@ def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=Fals """ cls._check_supports_udf("register_anonymous") if parameterized: + _validate_ret_dtype(ret_dtype, "unary", is_udt=is_udt, parameterized=True) return ParameterizedUnaryOp(name, func, anonymous=True, is_udt=is_udt) - return cls._build(name, func, anonymous=True, is_udt=is_udt) + return cls._build(name, func, anonymous=True, is_udt=is_udt, ret_dtype=ret_dtype) @classmethod - def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=False): + def register_new( + cls, name, func, *, parameterized=False, is_udt=False, lazy=False, ret_dtype=None + ): """Register a new UnaryOp and save it to ``graphblas.unary`` namespace. Parameters @@ -360,6 +384,9 @@ def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=Fal Compiling functions can be slow, however, so you may want to delay compilation and only compile when the operator is used, which is done by setting ``lazy=True``. + ret_dtype : dtype, optional + The dtype the operator returns. Requires ``is_udt=True``. + See :meth:`register_anonymous` for details. Examples -------- @@ -369,17 +396,27 @@ def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=Fal """ cls._check_supports_udf("register_new") + # Validate eagerly even for lazy=True, so a bad combination fails at + # the registration site rather than at first attribute touch. + _validate_ret_dtype(ret_dtype, "unary", is_udt=is_udt, parameterized=parameterized) module, funcname = cls._remove_nesting(name) if lazy: module._delayed[funcname] = ( cls.register_new, - {"name": name, "func": func, "parameterized": parameterized, "is_udt": is_udt}, + { + "name": name, + "func": func, + "parameterized": parameterized, + "is_udt": is_udt, + "ret_dtype": ret_dtype, + }, ) elif parameterized: + _validate_ret_dtype(ret_dtype, "unary", is_udt=is_udt, parameterized=True) unary_op = ParameterizedUnaryOp(name, func, is_udt=is_udt) setattr(module, funcname, unary_op) else: - unary_op = cls._build(name, func, is_udt=is_udt) + unary_op = cls._build(name, func, is_udt=is_udt, ret_dtype=ret_dtype) setattr(module, funcname, unary_op) # Also save it to `graphblas.op` if not yet defined opmodule, funcname = cls._remove_nesting(name, module=op, modname="op", strict=False) @@ -453,7 +490,7 @@ def _initialize(cls): ]: unop.orig_func = func if _has_numba: - unop._numba_func = numba.njit(func) + unop._numba_func = numba.njit(func, error_model="numpy") else: unop._numba_func = None unop._udt_types = {} @@ -476,12 +513,14 @@ def __init__( is_positional=False, is_udt=False, numba_func=None, + ret_dtype=None, ): super().__init__(name, anonymous=anonymous) self.orig_func = func self._numba_func = numba_func self.is_positional = is_positional self._is_udt = is_udt + self._ret_dtype = ret_dtype if is_udt: self._udt_types = {} # {dtype: DataType} self._udt_ops = {} # {dtype: TypedUserUnaryOp} diff --git a/graphblas/core/operator/utils.py b/graphblas/core/operator/utils.py index 6f2df5535..b293c2115 100644 --- a/graphblas/core/operator/utils.py +++ b/graphblas/core/operator/utils.py @@ -57,6 +57,21 @@ raise +# ``agg.py`` imports ``get_typed_op`` from this module, so the two form an import +# cycle and we can't import the Aggregator classes at module load. ``agg.py`` +# registers them here when it loads, which ``core.operator.__init__`` does +# eagerly at import. Before registration no Aggregator instance can exist, so +# ``get_typed_op`` skips the check entirely. +_Aggregator = None +_TypedAggregator = None + + +def _register_aggregator_types(aggregator, typed_aggregator): + global _Aggregator, _TypedAggregator + _Aggregator = aggregator + _TypedAggregator = typed_aggregator + + def get_typed_op(op, dtype, dtype2=None, *, is_left_scalar=False, is_right_scalar=False, kind=None): if isinstance(op, OpBase): # UDTs always get compiled @@ -93,15 +108,14 @@ def get_typed_op(op, dtype, dtype2=None, *, is_left_scalar=False, is_right_scala if isinstance(op, TypedOpBase): return op - from .agg import Aggregator, TypedAggregator - - if isinstance(op, Aggregator): - # agg._any_dtype basically serves the same purpose as op._custom_dtype - if op._any_dtype is not None and op._any_dtype is not True: - return op[op._any_dtype] - return op[dtype] - if isinstance(op, TypedAggregator): - return op + if _Aggregator is not None: + if isinstance(op, _Aggregator): + # agg._any_dtype basically serves the same purpose as op._custom_dtype + if op._any_dtype is not None and op._any_dtype is not True: + return op[op._any_dtype] + return op[dtype] + if isinstance(op, _TypedAggregator): + return op if isinstance(op, str): if kind == "unary": op = unary_from_string(op) @@ -280,6 +294,52 @@ def get_semiring(monoid, binaryop, name=None): return rv +def _resolve_index_expr(ns, modname, expr, callname, opname): + """Turn an infix comparison expression into an indexunary/select op call. + + Shared by the ``value``/``row``/``column`` helpers in the ``select`` and + ``indexunary`` namespaces (and ``select.index``). ``select.valuegt`` is a + SelectOp that defaults to ``.select``, while ``indexunary.valuegt`` is an + IndexUnaryOp that defaults to ``.apply``, so the same expression resolves to + a select or an apply depending on which namespace's ops ``ns`` holds. ``ns`` + is the namespace module's ``globals()`` dict and ``modname`` is its short + name (used in the error messages). + """ + from ..base import BaseExpression + + if not isinstance(expr, BaseExpression): + raise TypeError( + f"Expected ScalarExpression, VectorExpression, or MatrixExpression; " + f"found {type(expr)}\nTypical usage: {modname}.{callname}(x <= 5)" + ) + tensor = expr.args[0] + thunk = expr.args[1] + method = f"{opname}{expr.op.name}" + if method not in ns: + # TODO: remove this once rowlt/rowge/collt/colge exist + # Convert thunk to Python int to avoid possible subtraction with uints + thunk = thunk.value + # Attempt to convert < into <= (rowlt is not part of official spec, but rowle is) + if expr.op.name == "lt": + method = f"{opname}le" + thunk -= 1 + # Attempt to convert >= into > (rowge is not part of official spec, but rowgt is) + elif expr.op.name == "ge": + method = f"{opname}gt" + thunk -= 1 + if method not in ns: # pragma: no cover (sanity) + raise ValueError(f"Unknown or unregistered {modname} method: {method}") + if expr._is_scalar: + # Handle ScalarExpressions that change their arguments to Vector + if tensor._parent is not None: # e.g., suitesparse + tensor = tensor._parent + thunk = thunk._parent + else: # e.g., suitesparse-vanilla + tensor = tensor[0].new() + thunk = thunk[0].new() + return ns[method](tensor, thunk) + + unary.register_new = UnaryOp.register_new unary.register_anonymous = UnaryOp.register_anonymous indexbinary.register_new = IndexBinaryOp.register_new @@ -417,28 +477,196 @@ def _from_string(string, module, mapping, example): def unary_from_string(string): + """Look up a UnaryOp by name or symbol, optionally typed with ``"[dtype]"``. + + Backs ``gb.unary.from_string`` and the string coercion used wherever a + UnaryOp is accepted, such as ``v.apply("abs")``. + + Parameters + ---------- + string : str + A name in the ``gb.unary`` namespace (``"abs"``, or a dotted path such + as ``"numpy.negative"``) or a symbolic shorthand (``"-"`` for ``ainv``, + ``"~"`` for ``lnot``). Append ``"[dtype]"`` to type the operator, as in + ``"abs[int]"``. + + Returns + ------- + UnaryOp + + See Also + -------- + unary.register_new + op.from_string + + Examples + -------- + >>> gb.unary.from_string("abs") is gb.unary.abs + True + >>> gb.unary.from_string("abs[int]") is gb.unary.abs[int] + True + + """ return _from_string(string, unary, _str_to_unary, "abs[int]") def indexunary_from_string(string): + """Look up an IndexUnaryOp by name, optionally typed with ``"[dtype]"``. + + Parameters + ---------- + string : str + A name in the ``gb.indexunary`` namespace, such as ``"rowindex"``, + ``"diag"``, or ``"tril"``. Append ``"[dtype]"`` to type the operator, + as in ``"rowindex[int]"``. + + Returns + ------- + IndexUnaryOp + + See Also + -------- + indexunary.register_new + select.from_string + + Examples + -------- + >>> gb.indexunary.from_string("rowindex") is gb.indexunary.rowindex + True + + """ # "select" is a variant of IndexUnary, so the string abbreviations in # _str_to_select are appropriate to reuse here - return _from_string(string, indexunary, _str_to_select, "row_index") + return _from_string(string, indexunary, _str_to_select, "rowindex") def select_from_string(string): + """Look up a SelectOp by name or comparison symbol. + + Parameters + ---------- + string : str + A name in the ``gb.select`` namespace (``"tril"``, ``"triu"``, + ``"offdiag"``, ``"valuegt"``, ...) or a comparison shorthand such as + ``">="`` (``valuege``), ``"=="`` (``valueeq``), or ``"row>"`` + (``rowgt``). + + Returns + ------- + SelectOp + + See Also + -------- + select.register_new + indexunary.from_string + + Examples + -------- + >>> gb.select.from_string("tril") is gb.select.tril + True + >>> gb.select.from_string(">=") is gb.select.valuege + True + + """ return _from_string(string, select, _str_to_select, "tril") def binary_from_string(string): + """Look up a BinaryOp by name or symbol, optionally typed with ``"[dtype]"``. + + Backs ``gb.binary.from_string`` and the string coercion used wherever a + BinaryOp is accepted, such as ``A.ewise_mult(B, "+")``. + + Parameters + ---------- + string : str + A name in the ``gb.binary`` namespace (``"plus"``, or a dotted path such + as ``"numpy.mod"``) or an arithmetic/comparison shorthand such as ``"+"`` + (``plus``), ``"*"`` (``times``), or ``">="`` (``ge``). Append + ``"[dtype]"`` to type the operator, as in ``"plus[int]"``. + + Returns + ------- + BinaryOp + + See Also + -------- + binary.register_new + op.from_string + + Examples + -------- + >>> gb.binary.from_string("+") is gb.binary.plus + True + >>> gb.binary.from_string("minus[int]") is gb.binary.minus[int] + True + + """ return _from_string(string, binary, _str_to_binary, "+[int]") def monoid_from_string(string): + """Look up a Monoid by name or symbol, optionally typed with ``"[dtype]"``. + + Parameters + ---------- + string : str + A name in the ``gb.monoid`` namespace (``"plus"``, ``"times"``, ...) or a + symbolic shorthand such as ``"+"`` (``plus``), ``"*"`` (``times``), or + ``"|"`` (``lor``). Append ``"[dtype]"`` to type the monoid, as in + ``"plus[float]"``. + + Returns + ------- + Monoid + + See Also + -------- + monoid.register_new + semiring.from_string + + Examples + -------- + >>> gb.monoid.from_string("+[float]") is gb.monoid.plus[float] + True + + """ return _from_string(string, monoid, _str_to_monoid, "+[int]") def semiring_from_string(string): + """Look up a Semiring by name, optionally typed with ``"[dtype]"``. + + A semiring pairs a monoid with a binaryop. Name it either as the combined + namespace attribute (``"plus_times"``) or in ``"monoid.binaryop"`` form using + the monoid and binaryop shorthands (``"min.+"``); the two parts must be + separated by exactly one period. + + Parameters + ---------- + string : str + The semiring name, such as ``"plus_times"`` or ``"min_plus"``, or the + ``"monoid.binaryop"`` form ``"min.+"``. Append ``"[dtype]"`` to type the + semiring, as in ``"min.+[int]"``. + + Returns + ------- + Semiring + + See Also + -------- + semiring.register_new + semiring.get_semiring + op.from_string + + Examples + -------- + >>> gb.semiring.from_string("min.+") is gb.semiring.min_plus + True + >>> gb.semiring.from_string("min_plus") is gb.semiring.min_plus + True + + """ split = string.split(".") if len(split) == 1: try: @@ -457,6 +685,38 @@ def semiring_from_string(string): def op_from_string(string): + """Look up an operator of any kind by string. + + Each operator type is tried in turn (unary, binary, monoid, semiring, + indexunary, select, then aggregator) and the first match is returned, so an + unqualified name resolves to whichever kind defines it first. Use a + type-specific ``from_string`` (e.g. ``gb.binary.from_string``) when the kind + is known and matters. + + Parameters + ---------- + string : str + An operator name or symbol accepted by any of the type-specific + ``from_string`` functions, optionally typed with ``"[dtype]"``. + + Returns + ------- + UnaryOp, BinaryOp, Monoid, Semiring, IndexUnaryOp, SelectOp, or Aggregator + + See Also + -------- + unary.from_string + binary.from_string + semiring.from_string + + Examples + -------- + >>> gb.op.from_string("+") is gb.binary.plus + True + >>> gb.op.from_string("min.plus") is gb.semiring.min_plus + True + + """ for func in [ # Note: order matters here unary_from_string, @@ -491,6 +751,32 @@ def op_from_string(string): def aggregator_from_string(string): + """Look up an Aggregator by name or symbol, optionally typed with ``"[dtype]"``. + + Parameters + ---------- + string : str + A name in the ``gb.agg`` namespace (``"sum"``, ``"count"``, ``"any"``, + ...) or a symbolic shorthand such as ``"+"`` (``sum``), ``"*"`` + (``prod``), ``"&"`` (``all``), or ``"|"`` (``any``). Append ``"[dtype]"`` + to type the aggregator, as in ``"sum[int]"``. + + Returns + ------- + Aggregator + + See Also + -------- + op.from_string + + Examples + -------- + >>> gb.agg.from_string("sum[int]") is gb.agg.sum[int] + True + >>> gb.agg.from_string("|") is gb.agg.any + True + + """ return _from_string(string, agg, _str_to_agg, "sum[int]") diff --git a/graphblas/core/scalar.py b/graphblas/core/scalar.py index 5b5e56299..4f0d7f317 100644 --- a/graphblas/core/scalar.py +++ b/graphblas/core/scalar.py @@ -7,7 +7,7 @@ from ..dtypes import _INDEX, FP64, _index_dtypes, lookup_dtype, unify from ..exceptions import EmptyObject, check_status from . import _has_numba, _supports_udfs, automethods, ffi, lib, utils -from .base import BaseExpression, BaseType, call +from .base import BaseExpression, BaseType, _is_recording, call from .expr import AmbiguousAssignOrExtract from .operator import get_typed_op from .utils import _Pointer, output_type, wrapdoc @@ -1082,7 +1082,37 @@ class ScalarIndexExpr(AmbiguousAssignOrExtract): def new(self, dtype=None, *, is_cscalar=None, name=None, **opts): if is_cscalar is None: is_cscalar = False - return self.parent._extract_element( + parent = self.parent + # Fast path for the default `expr.new()`: extract a single element + # straight into a fresh GrB_Scalar via GrB_*_extractElement_Scalar, + # skipping the `call` wrapper's per-arg _carg marshalling. Falls back + # for a dtype cast, cscalar output, opts, UDTs, and an active Recorder + # (so the call is recorded). Result is a GrB_Scalar (is_cscalar=False), + # empty exactly when the element is missing, same as _extract_element. + if ( + dtype is None + and not is_cscalar + and not opts + and not parent.dtype._is_udt + and not _is_recording() + ): + indices = self.resolved_indexes.indices + result = Scalar(parent.dtype, is_cscalar=False, name=name) # pragma: is_grbscalar + if len(indices) == 1: + err_code = lib.GrB_Vector_extractElement_Scalar( + result.gb_obj[0], parent.gb_obj[0], indices[0].index._carg + ) + else: + rowidx, colidx = indices + if parent._is_transposed: + rowidx, colidx = colidx, rowidx + err_code = lib.GrB_Matrix_extractElement_Scalar( + result.gb_obj[0], parent.gb_obj[0], rowidx.index._carg, colidx.index._carg + ) + if err_code: + check_status(err_code, [result]) + return result + return parent._extract_element( self.resolved_indexes, dtype, opts, is_cscalar=is_cscalar, name=name ) @@ -1094,6 +1124,21 @@ def dup(self, dtype=None, *, clear=False, is_cscalar=False, name=None, **opts): return Scalar(dtype, is_cscalar=is_cscalar, name=name) return self.new(dtype, is_cscalar=is_cscalar, name=name, **opts) + def _extract_fast(self): + """Resolve a value read (``.value``, ``float(...)``, ...) with one extract. + + Those readers only need the raw element, so extract it straight into a + cscalar and skip the extra GrB_Scalar round-trip that ``.new()`` followed + by ``Scalar.value`` would perform. Defer to the full ``.new()`` for UDTs + (whose values need numpy conversion in ``Scalar.value``) and while a + Recorder is active (so it observes the same calls as the expression path). + ``automethods._get_value`` consults this hook for ``_fast_scalar_attrs``. + """ + parent = self.parent + if parent.dtype._is_udt or _is_recording(): + return self.new() + return parent._extract_element(self.resolved_indexes, None, {}, is_cscalar=True) + is_cscalar = Scalar.is_cscalar is_grbscalar = Scalar.is_grbscalar __hash__ = None diff --git a/graphblas/core/ss/config.py b/graphblas/core/ss/config.py index 70a7dd196..4d2a20573 100644 --- a/graphblas/core/ss/config.py +++ b/graphblas/core/ss/config.py @@ -52,8 +52,38 @@ def __init__(self, parent=None, context=None): if k not in rd: # pragma: no branch (safety) rd[k] = k cls._initialized = True + # Writes are unrestricted until __init__ finishes, and the instance + # attributes set by then become the allowlist for __setattr__. + # Subclasses that add instance attributes (such as Context) must set + # them before calling super().__init__(). + object.__setattr__(self, "_initializing", True) self._parent = parent self._context = context + object.__setattr__(self, "_internal_attrs", frozenset(self.__dict__) - {"_initializing"}) + object.__setattr__(self, "_initializing", False) + + def __setattr__(self, key, value): + # An instance built without running __init__ has no allowlist yet, so stay + # permissive rather than raising a confusing AttributeError from the guard. + # Properties (such as Context._context) go through so their setters run. + if ( + self.__dict__.get("_initializing", True) + or key in self._internal_attrs + or isinstance(getattr(type(self), key, None), property) + ): + object.__setattr__(self, key, value) + return + if (option := key.lower()) in self._options: + if option in self._read_only: + raise AttributeError(f"Config option {option!r} is read-only") + raise AttributeError( + f"Cannot set config option {option!r} by attribute assignment; " + f"this does not change the config. " + f"Use item assignment instead: config[{option!r}] = value" + ) + raise AttributeError( + f"Unknown config option {key!r}; known options are {sorted(self._options)}." + ) def __delitem__(self, key): raise TypeError("Configuration options can't be deleted.") @@ -210,5 +240,5 @@ def __repr__(self): + "})" ) - def _ipython_key_completions_(self): # pragma: no cover (ipython) + def _ipython_key_completions_(self): return list(self) diff --git a/graphblas/core/ss/context.py b/graphblas/core/ss/context.py index 67c7a6c20..99026e16f 100644 --- a/graphblas/core/ss/context.py +++ b/graphblas/core/ss/context.py @@ -33,9 +33,12 @@ class Context(BaseConfig): } def __init__(self, engage=True, *, stack=True, nthreads=None, chunk=None, gpu_id=None): - super().__init__() + # Instance attributes must exist before super().__init__() finishes, + # because BaseConfig.__setattr__ allows only the attributes set by then. self.gb_obj = ffi_new("GxB_Context*") check_status_carg(lib.GxB_Context_new(self.gb_obj), "Context", self.gb_obj[0]) + self._prev_context = None + super().__init__() if stack: context = threadlocal.context self["nthreads"] = context["nthreads"] if nthreads is None else nthreads @@ -49,7 +52,6 @@ def __init__(self, engage=True, *, stack=True, nthreads=None, chunk=None, gpu_id self["chunk"] = chunk if gpu_id is not None and "gpu_id" in self._options: self["gpu_id"] = gpu_id - self._prev_context = None if engage: self.engage() diff --git a/graphblas/core/utils.py b/graphblas/core/utils.py index 49a8bd79c..163c40e27 100644 --- a/graphblas/core/utils.py +++ b/graphblas/core/utils.py @@ -33,6 +33,23 @@ def inner(func_wo_doc): return inner +def _module_attr_error(module_name, key, names): + """Build the AttributeError raised by an operator namespace's ``__getattr__``. + + ``names`` should be the module's ``__dir__()`` so lazily-registered operators + are offered as "did you mean" suggestions without forcing them to build. + """ + import difflib + + msg = f"module {module_name!r} has no attribute {key!r}" + candidates = [name for name in names if not name.startswith("_")] + matches = difflib.get_close_matches(key, candidates, n=3) + if matches: + hint = " or ".join(repr(match) for match in matches) + msg = f"{msg}. Did you mean {hint}?" + return AttributeError(msg) + + # Include most common types (even mistakes) _output_types = { int: int, diff --git a/graphblas/core/vector.py b/graphblas/core/vector.py index 8c73ecc48..ef672175e 100644 --- a/graphblas/core/vector.py +++ b/graphblas/core/vector.py @@ -4,12 +4,13 @@ from .. import backend, binary, monoid, select, semiring, unary from ..dtypes import _INDEX, FP64, INT64, lookup_dtype, unify -from ..exceptions import DimensionMismatch, NoValue, check_status +from ..exceptions import DimensionMismatch, GrB_NO_VALUE, NoValue, check_status, check_status_carg from . import _supports_udfs, automethods, ffi, lib, utils -from .base import BaseExpression, BaseType, _check_mask, call +from .base import BaseExpression, BaseType, _is_recording, call from .descriptor import lookup as descriptor_lookup +from .dtypes import _raise_dtype_or_arraylike from .expr import _ALL_INDICES, AmbiguousAssignOrExtract, IndexerResolver, InfixExprBase, Updater -from .mask import Mask, StructuralMask, ValueMask +from .mask import Mask, StructuralMask, ValueMask, _check_mask from .operator import ( UNKNOWN_OPCLASS, _get_typed_op_from_exprs, @@ -152,21 +153,22 @@ class Vector(BaseType): """ - __slots__ = "_size", "_parent", "ss" + __slots__ = "_size", "_parent" ndim = 1 _name_counter = itertools.count() def __new__(cls, dtype=FP64, size=0, *, name=None): self = object.__new__(cls) - self.dtype = lookup_dtype(dtype) + try: + self.dtype = lookup_dtype(dtype) + except (ValueError, TypeError) as exc: + _raise_dtype_or_arraylike("Vector", dtype, exc) size = _as_scalar(size, _INDEX, is_cscalar=True) self.name = f"v_{next(Vector._name_counter)}" if name is None else name self.gb_obj = ffi_new("GrB_Vector*") call("GrB_Vector_new", [_Pointer(self), self.dtype, size]) self._size = size.value self._parent = None - if backend == "suitesparse": - self.ss = ss(self) return self @classmethod @@ -177,8 +179,6 @@ def _from_obj(cls, gb_obj, dtype, size, *, parent=None, name=None): self.dtype = dtype self._size = size self._parent = parent - if backend == "suitesparse": - self.ss = ss(self) return self def __del__(self): @@ -303,6 +303,39 @@ def __setitem__(self, keys, expr, **opts): v[:] = 1 """ + # Fast path for `v[i] = scalar`: a plain integer index and an exact-fit + # Python scalar, with no mask/accum/opts, a non-UDT dtype, and no active + # Recorder. This mirrors what Updater -> _assign_element does for a + # single element, but skips building the resolver, Updater, and Scalar. + # Only int/float/bool/complex values are taken here so the dtype + # inference and cffi coercion match _assign_element exactly, and only + # int and np.integer keys (bools excluded) so the accepted indices + # match parse_index; everything else (slices, fancy indexing, 0-d + # arrays and other __index__ objects, numpy or Scalar values, + # `v(mask)[i] << x`) falls back to the full assign path, leaving + # mask/accum, coercion, and index errors unchanged. + if ( + not opts + and type(expr) in (int, float, bool, complex) + and (type(keys) is int or isinstance(keys, np.integer)) + and not self.dtype._is_udt + and not _is_recording() + ): + idx = keys.__index__() + size = self._size + if idx < 0: + idx += size + if idx < 0 or idx >= size: + raise IndexError(f"Index out of range: index={keys}, size={size}") + vdtype = lookup_dtype(type(expr), expr) + cvalue = ffi_new(f"{vdtype.c_type}*") + cvalue[0] = expr # cffi coercion, identical to the Scalar.value setter + err_code = utils.libget(f"GrB_Vector_setElement_{vdtype.name}")( + self.gb_obj[0], cvalue[0], idx + ) + if err_code: + check_status_carg(err_code, "Vector", self.gb_obj[0]) + return Updater(self, opts=opts)[keys] = expr def __contains__(self, index): @@ -316,6 +349,34 @@ def __contains__(self, index): 15 in v """ + # Fast path for a plain integer index: probe with + # GrB_Vector_extractElement directly instead of building an extract + # expression and Scalar. An out-of-range index falls through to the + # expression path so it raises the same IndexError as the slow path. + # Only int and np.integer take the fast lane (bools excluded), matching + # parse_index; other __index__ objects fall through to the expression + # path and its canonical errors. Fall back for a UDT dtype and an active + # Recorder (so the call is recorded), mirroring Vector.get. + if ( + (type(index) is int or isinstance(index, np.integer)) + and not self.dtype._is_udt + and not _is_recording() + ): + idx = index.__index__() + size = self._size + if idx < 0: + idx += size + if 0 <= idx < size: + dtype = self.dtype + res = ffi_new(f"{dtype.c_type}*") + err_code = utils.libget(f"GrB_Vector_extractElement_{dtype.name}")( + res, self.gb_obj[0], idx + ) + if err_code: + if err_code == GrB_NO_VALUE: + return False + check_status_carg(err_code, "Vector", self.gb_obj[0]) + return True extractor = self[index] if not extractor._is_scalar: raise TypeError( @@ -669,6 +730,35 @@ def get(self, index, default=None): Python scalar """ + # Fast path for a plain integer index: call GrB_Vector_extractElement + # directly instead of building an extract expression, which costs ~10x + # more than the C call for single-element access. Fall back when a + # Recorder is active (so the call is recorded) and for UDTs (whose + # values need numpy-based conversion in Scalar.value). Only int and + # np.integer take the fast lane (bools excluded), matching parse_index; + # other __index__ objects fall through to the expression path and its + # canonical errors. + if ( + (type(index) is int or isinstance(index, np.integer)) + and not self.dtype._is_udt + and not _is_recording() + ): + idx = index.__index__() + size = self._size + if idx < 0: + idx += size + if idx < 0 or idx >= size: + raise IndexError(f"Index out of range: index={index}, size={size}") + dtype = self.dtype + res = ffi_new(f"{dtype.c_type}*") + err_code = utils.libget(f"GrB_Vector_extractElement_{dtype.name}")( + res, self.gb_obj[0], idx + ) + if err_code: + if err_code == GrB_NO_VALUE: + return default + check_status_carg(err_code, "Vector", self.gb_obj[0]) + return res[0] expr = self[index] if expr._is_scalar: rv = expr.new().value @@ -682,6 +772,11 @@ def get(self, index, default=None): def from_coo(cls, indices, values=1.0, dtype=None, *, size=None, dup_op=None, name=None): """Create a new Vector from indices and values. + .. warning:: + When ``size`` is omitted, it is inferred from the largest index, so + trailing empty positions are dropped. Pass ``size`` explicitly to + pin the length. + Parameters ---------- indices : list or np.ndarray @@ -1418,7 +1513,7 @@ def apply(self, op, right=None, *, left=None): """ method_name = "apply" extra_message = ( - "apply only accepts UnaryOp with no scalars or BinaryOp with `left` or `right` scalar" + "apply only accepts UnaryOp with no scalars or BinaryOp with `left` or `right` scalar " "or IndexUnaryOp with `right` thunk." ) if isinstance(op, str): @@ -2082,10 +2177,7 @@ def from_dict(cls, d, dtype=None, *, size=None, name=None): else: # If we know the dtype, then using `np.fromiter` is much faster dtype = lookup_dtype(dtype) - if dtype.np_type.subdtype is not None and np.__version__[:5] in {"1.21.", "1.22."}: - values, dtype = values_to_numpy_buffer(list(d.values()), dtype) # FLAKY COVERAGE - else: - values = np.fromiter(d.values(), dtype.np_type) + values = np.fromiter(d.values(), dtype.np_type) if size is None and indices.size == 0: size = 0 return cls.from_coo(indices, values, dtype, size=size, name=name) @@ -2109,10 +2201,17 @@ def to_dict(self): if backend == "suitesparse": - Vector.ss = class_property(Vector.ss, ss) + # Built lazily per access, not stored, to avoid the ss/_parent reference + # cycle that would keep every Vector alive until the cyclic gc; see gh-559 + # and the matching note in matrix.py. + def _ss(self): + return ss(self) + + _ss.__name__ = _ss.__qualname__ = "ss" + Vector.ss = class_property(property(_ss), ss) else: Vector.ss = class_property( - Vector.ss, 'ss attribute is only available with "suitesparse" backend', exceptional=True + property(), 'ss attribute is only available with "suitesparse" backend', exceptional=True ) diff --git a/graphblas/indexunary/__init__.py b/graphblas/indexunary/__init__.py index a3cb06608..4aee3524a 100644 --- a/graphblas/indexunary/__init__.py +++ b/graphblas/indexunary/__init__.py @@ -25,7 +25,60 @@ def __getattr__(key): ss = import_module(".ss", __name__) globals()["ss"] = ss return ss - raise AttributeError(f"module {__name__!r} has no attribute {key!r}") + from ..core.utils import _module_attr_error + + raise _module_attr_error(__name__, key, __dir__()) + + +def _resolve_expr(expr, callname, opname): + from ..core.operator.utils import _resolve_index_expr + + return _resolve_index_expr(globals(), "indexunary", expr, callname, opname) + + +def value(expr): + """An advanced indexunary method for easily expressing value comparison logic. + + Example usage: + >>> gb.indexunary.value(A > 0) + + The example will dispatch to ``gb.indexunary.valuegt(A, 0)`` + while being nicer to read. + """ + return _resolve_expr(expr, "value", "value") + + +def row(expr): + """An advanced indexunary method for easily expressing Matrix row index comparison logic. + + Example usage: + >>> gb.indexunary.row(A <= 5) + + The example will dispatch to ``gb.indexunary.rowle(A, 5)`` + while being potentially nicer to read. + """ + return _resolve_expr(expr, "row", "row") + + +def column(expr): + """An advanced indexunary method for easily expressing Matrix column index comparison logic. + + Example usage: + >>> gb.indexunary.column(A <= 5) + + The example will dispatch to ``gb.indexunary.colle(A, 5)`` + while being potentially nicer to read. + """ + return _resolve_expr(expr, "column", "col") + + +# Note: an ``index`` helper (the Vector analogue of ``select.index``) is *not* +# provided here because ``indexunary.index`` already exists as an alias for the +# positional ``rowindex`` op (INT64). It is relied on as an operator, e.g. +# ``v.apply(indexunary.index)``, which a helper function would break. For a +# Vector index comparison use ``indexunary.row(v < k)`` (resolves to ``rowle``, +# the same op ``select.index`` uses) or the explicit ``indexunary.indexle`` / +# ``indexunary.indexgt``. from ..core import operator # noqa: E402 isort:skip diff --git a/graphblas/io/_networkx.py b/graphblas/io/_networkx.py index 8cf84e576..c4a31e6db 100644 --- a/graphblas/io/_networkx.py +++ b/graphblas/io/_networkx.py @@ -27,6 +27,121 @@ def from_networkx(G, nodelist=None, dtype=None, weight="weight", name=None): if dtype is not None: dtype = lookup_dtype(dtype).np_type + + # The node selection below mirrors nx.to_scipy_sparse_array so that empty + # graphs, nodelist subsets, missing nodes, and duplicate nodes raise the + # same errors and produce the same ordering. Building the coo arrays here + # (instead of via a scipy sparse round-trip) skips the extra coo -> csr + # materialization and the scipy.sparse import, which alone costs over 100ms. + import numpy as np + + from ..binary import plus + from ..core.matrix import Matrix + + if len(G) == 0: + raise nx.NetworkXError("Graph has no nodes or edges") + + if nodelist is None: + nodelist = list(G) + nlen = len(G) + else: + nlen = len(nodelist) + if nlen == 0: + raise nx.NetworkXError("nodelist has no nodes") + nodeset = set(G.nbunch_iter(nodelist)) + if nlen != len(nodeset): + for n in nodelist: + if n not in G: + raise nx.NetworkXError(f"Node {n} in nodelist is not in G") + raise nx.NetworkXError("nodelist contains duplicates.") + if nlen < len(G): + G = G.subgraph(nodelist) + + index = dict(zip(nodelist, range(nlen), strict=True)) + coefficients = zip( + *((index[u], index[v], wt) for u, v, wt in G.edges(data=weight, default=1)), + strict=True, + ) + try: + row, col, data = coefficients + except ValueError: + # there is no edge in the (sub)graph + row, col, data = (), (), () + + if G.is_directed(): + rows, cols, vals = row, col, data + # A multigraph can have parallel edges (duplicate ``(u, v)``); summing + # them with ``plus`` matches scipy's coo -> csr accumulation. A simple + # graph has no duplicate coordinates, so ``dup_op=None`` is exact and + # skips the accumulator. + dup_op = plus if G.is_multigraph() else None + else: + # Symmetrize: mirror off-diagonal entries. Self-loops would be double + # counted, so subtract the diagonal contribution once, matching + # nx.to_scipy_sparse_array. dup_op=plus then sums the diagonal triple + # (wt + wt - wt) back to wt. For a multigraph, plus also sums parallel + # edges (duplicate coordinates) the same way scipy's coo -> csr does; + # for a simple graph off-diagonal entries are unique so plus is a no-op + # there. + d = data + data + r = row + col + c = col + row + selfloops = list(nx.selfloop_edges(G, data=weight, default=1)) + if selfloops: + diag_index, diag_data = zip(*((index[u], -wt) for u, v, wt in selfloops), strict=True) + d += diag_data + r += diag_index + c += diag_index + rows, cols, vals = r, c, d + dup_op = plus + + values = np.array(vals, dtype=dtype) + if dtype is None and values.dtype == np.int32: # pragma: no cover (win64 numpy < 2) + # numpy < 2 infers the platform C long for a sequence of Python ints, which + # is 32-bit on Windows. values_to_numpy_buffer widens the same way for + # non-numpy input, so this keeps from_networkx agreeing with + # Matrix.from_coo on INT64 for an unweighted graph on every platform. + values = values.astype(np.int64) + if values.ndim != 1 or values.dtype.kind not in "biufc": + # Defer to scipy so the error matches the previous behavior exactly. + # Two kinds of weight land here: non-numeric attributes (object arrays, + # but also e.g. all-string weights, which infer a csr conversion instead: scipy raises + # TypeError ("no supported conversion for types"), which networkx 3.4+ + # wraps in a NetworkXError that blames the sparse format while + # networkx <= 3.3 lets the TypeError propagate. + # Restate it so every supported stack reports the same error for the + # same graph. The graph and nodelist checks above are the ones + # nx.to_scipy_sparse_array makes, so a NetworkXError or TypeError from + # the fallback can only be that dtype complaint. + raise ValueError( + f"scipy.sparse does not support dtype {values.dtype}; " + "edge weights must be numeric scalars" + ) from err + if values.size == 0: + # An empty graph has no data to infer a dtype from; scipy defaults an + # empty coo array to float64, so match that when dtype is unset. + return Matrix(lookup_dtype(values.dtype), nrows=nlen, ncols=nlen, name=name) + + rows = np.array(rows, dtype=np.uint64) + cols = np.array(cols, dtype=np.uint64) + return Matrix.from_coo(rows, cols, values, nrows=nlen, ncols=nlen, dup_op=dup_op, name=name) + + +def _from_networkx_via_scipy(G, nodelist, dtype, weight, name): + """Fallback path: convert through a scipy sparse array. + + ``dtype`` is already normalized to a numpy type (or None) by the caller. + """ + import networkx as nx + A = nx.to_scipy_sparse_array(G, nodelist=nodelist, dtype=dtype, weight=weight) return from_scipy_sparse(A, name=name) diff --git a/graphblas/monoid/__init__.py b/graphblas/monoid/__init__.py index 027fc0afe..30463e6a2 100644 --- a/graphblas/monoid/__init__.py +++ b/graphblas/monoid/__init__.py @@ -29,7 +29,9 @@ def __getattr__(key): ss = import_module(".ss", __name__) globals()["ss"] = ss return ss - raise AttributeError(f"module {__name__!r} has no attribute {key!r}") + from ..core.utils import _module_attr_error + + raise _module_attr_error(__name__, key, __dir__()) from ..core import operator # noqa: E402 isort:skip diff --git a/graphblas/op/__init__.py b/graphblas/op/__init__.py index 1eb2b51d7..5a37b6121 100644 --- a/graphblas/op/__init__.py +++ b/graphblas/op/__init__.py @@ -47,7 +47,9 @@ def __getattr__(key): f"module {__name__!r} unable to compile UDF for {key!r}; " "install numba for UDF support" ) - raise AttributeError(f"module {__name__!r} has no attribute {key!r}") + from ..core.utils import _module_attr_error + + raise _module_attr_error(__name__, key, __dir__()) from ..core import operator, _supports_udfs # noqa: E402 isort:skip diff --git a/graphblas/select/__init__.py b/graphblas/select/__init__.py index b55766ff8..8d9e5a3c5 100644 --- a/graphblas/select/__init__.py +++ b/graphblas/select/__init__.py @@ -29,43 +29,15 @@ def __getattr__(key): ss = import_module(".ss", __name__) globals()["ss"] = ss return ss - raise AttributeError(f"module {__name__!r} has no attribute {key!r}") + from ..core.utils import _module_attr_error + + raise _module_attr_error(__name__, key, __dir__()) def _resolve_expr(expr, callname, opname): - from ..core.base import BaseExpression - - if not isinstance(expr, BaseExpression): - raise TypeError( - f"Expected ScalarExpression, VectorExpression, or MatrixExpression; " - f"found {type(expr)}\nTypical usage: select.{callname}(x <= 5)" - ) - tensor = expr.args[0] - thunk = expr.args[1] - method = f"{opname}{expr.op.name}" - if method not in globals(): - # TODO: remove this once rowlt/rowge/collt/colge exist - # Convert thunk to Python int to avoid possible subtraction with uints - thunk = thunk.value - # Attempt to convert < into <= (rowlt is not part of official spec, but rowle is) - if expr.op.name == "lt": - method = f"{opname}le" - thunk -= 1 - # Attempt to convert >= into > (rowge is not part of official spec, but rowgt is) - elif expr.op.name == "ge": - method = f"{opname}gt" - thunk -= 1 - if method not in globals(): # pragma: no cover (sanity) - raise ValueError(f"Unknown or unregistered select method: {method}") - if expr._is_scalar: - # Handle ScalarExpressions that change their arguments to Vector - if tensor._parent is not None: # e.g., suitesparse - tensor = tensor._parent - thunk = thunk._parent - else: # e.g., suitesparse-vanilla - tensor = tensor[0].new() - thunk = thunk[0].new() - return globals()[method](tensor, thunk) + from ..core.operator.utils import _resolve_index_expr + + return _resolve_index_expr(globals(), "select", expr, callname, opname) def _match_expr(parent, expr): diff --git a/graphblas/semiring/__init__.py b/graphblas/semiring/__init__.py index 95a44261a..b0cf07bce 100644 --- a/graphblas/semiring/__init__.py +++ b/graphblas/semiring/__init__.py @@ -74,7 +74,9 @@ def __getattr__(key): f"module {__name__!r} unable to compile UDF for {key!r}; " "install numba for UDF support" ) - raise AttributeError(f"module {__name__!r} has no attribute {key!r}") + from ..core.utils import _module_attr_error + + raise _module_attr_error(__name__, key, __dir__()) from ..core import operator # noqa: E402 isort:skip diff --git a/graphblas/ss/_core.py b/graphblas/ss/_core.py index 91e1496e9..1c005f294 100644 --- a/graphblas/ss/_core.py +++ b/graphblas/ss/_core.py @@ -316,6 +316,13 @@ def __getitem__(self, key): raise KeyError(key) raise _error_code_lookup[info](f"Failed to get info for {key}") # pragma: no cover (safety) + def __setattr__(self, key, value): + # About never sets instance attributes, so any attribute write is a + # mistake that would otherwise create a dead attribute shadowing reads. + if key.lower() in self: + raise AttributeError(f"About option {key.lower()!r} is read-only") + raise AttributeError(f"Unknown About option {key!r}; known options are {sorted(self)}.") + def __iter__(self): return iter( sorted( diff --git a/graphblas/tests/test_autogenerate_check.py b/graphblas/tests/test_autogenerate_check.py new file mode 100644 index 000000000..a10beb561 --- /dev/null +++ b/graphblas/tests/test_autogenerate_check.py @@ -0,0 +1,62 @@ +"""Guard the auto-generated code blocks against undetected drift. + +``test_automethods.py`` guards the *names* on the generated expression surface. +Nothing there guards the generated *content*: hand-edit a block between the +``# Begin auto-generated code`` / ``# End auto-generated code`` markers, or edit +the name sets in ``automethods._main`` without rerunning the generator, and the +on-disk blocks go stale while CI stays green. + +``scripts/autogenerate.py --check`` closes that gap. It regenerates every +generated file into a scratch tree and compares each against the file on disk, +exiting nonzero (and naming the drifted files) on any mismatch. The comparison is +on parsed syntax, so the check needs no formatter and runs everywhere rather than +skipping in every job that lacks one. + +Scope, since the two guards are easy to confuse. The generator emits from the +literal name sets in ``automethods._main``, not by introspecting the classes, so +adding a method to Matrix/Vector/Scalar does not change its output and cannot +show up here. That case is ``test_automethods.py``'s: it compares the concrete +types against the generated surface and fails when a new public attribute is +neither auto-computed nor opted out. This module covers the other half, the +generated files no longer matching the generator that claims to produce them. +""" + +import subprocess +import sys +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +AUTOGEN_SCRIPT = REPO_ROOT / "scripts" / "autogenerate.py" +PACKAGE_LINE = "graphblas package: " + + +def test_autogenerated_code_is_in_sync(): + if not AUTOGEN_SCRIPT.exists(): # pragma: no cover (source checkout only) + pytest.skip("scripts/autogenerate.py not present (installed package, not a checkout)") + result = subprocess.run( + [sys.executable, str(AUTOGEN_SCRIPT), "--check"], + capture_output=True, + text=True, + check=False, + ) + report = f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + checked = [ + line[len(PACKAGE_LINE) :].strip() + for line in result.stdout.splitlines() + if line.startswith(PACKAGE_LINE) + ] + + # Assert the positive post-condition before the exit code. A check that ran against a + # different `graphblas` reports success about a tree nobody asked about: sys.path[0] + # for a script is the script's own directory, so that is what happens whenever an + # editable install points elsewhere, as in a git worktree. + assert checked, f"--check never reported a package, so it failed before starting.\n{report}" + expected = str(REPO_ROOT / "graphblas") + assert checked == [expected], f"--check validated {checked}, not {expected}.\n{report}" + + assert result.returncode == 0, ( + "Auto-generated code is out of date. Run `python scripts/autogenerate.py` and " + f"commit the result.\n{report}" + ) diff --git a/graphblas/tests/test_automethods.py b/graphblas/tests/test_automethods.py new file mode 100644 index 000000000..8cb2c25be --- /dev/null +++ b/graphblas/tests/test_automethods.py @@ -0,0 +1,372 @@ +"""Guard the auto-generated expression surface against silent drift. + +``graphblas/core/automethods.py`` is a generated-code module. Its name sets +(near line 347) drive ``scripts/autogenerate.py``, which copies auto-compute +properties onto the expression classes (``VectorExpression`` / +``VectorIndexExpr`` / ``VectorInfixExpr`` and the Scalar/Matrix equivalents) so +that, for example, ``(A @ B).to_coo()`` works without a manual ``.new()`` first. +The same generator run emits the surface onto the infix classes too (the +infix.py branch of ``automethods._main()``), minus the private ``_get_value`` +helper, so the infix classes must carry the identical public surface. + +The trap this module closes: add a public method to ``Matrix``/``Vector``/ +``Scalar``, forget to add its name to the sets and rerun the generator, and +nothing fails. The expression classes silently lack the method, but the concrete +types have it, so CI stays green. A name that lands on the plain expression +classes but is missing from the infix classes slips through the same way. + +The tests here assert, for each concrete type and for ``TransposedMatrix``: + +1. Forward: every public method/property is EITHER reachable on ALL of the + expression classes (plain, index, and infix) via auto-compute OR listed in + ``OPT_OUT`` with a reason. Mutating methods, constructors, and cheap metadata + deliberately do not auto-compute. +2. Reverse: every auto-generated name still exists on the concrete type (a + rename that leaves a stale set entry is caught at import already, but this + makes the failure legible). +3. Hygiene: every ``OPT_OUT`` entry is a live attribute that is not in fact + covered, so the table cannot rot into stale excuses. + +Coverage is derived at runtime from what the generator actually emitted onto the +expression classes (properties whose getter lives in the ``automethods`` module), +intersected across every variant so a name reachable on some classes but absent +from another (e.g. only the infix class) counts as uncovered. It is not a second +copy of the name sets, which keeps this test honest if the sets are reorganized: +only a real change in the generated surface moves coverage. + +Dunder scope: value-forwarding dunders such as ``__getitem__``, ``__contains__``, +and ``__matmul__`` are generated and thus checked. Arithmetic and comparison +operator sugar such as ``__add__`` and ``__lt__`` is handled by the infix +expression system (``infixmethods.py``), not the value-forwarding automethods +path, so it is excluded from the dunder sweep by ``_OPERATOR_SUGAR_DUNDERS``. +Python object machinery is excluded via a baseline class so new +interpreter-version dunders never make this test flaky. +""" + +from graphblas.core.infix import MatrixInfixExpr, ScalarInfixExpr, VectorInfixExpr +from graphblas.core.matrix import ( + Matrix, + MatrixExpression, + MatrixIndexExpr, + TransposedMatrix, +) +from graphblas.core.scalar import Scalar, ScalarExpression, ScalarIndexExpr +from graphblas.core.vector import Vector, VectorExpression, VectorIndexExpr + +import pytest # isort: skip + +_AUTOMETHODS_MODULE = "graphblas.core.automethods" + +# --- reasons an attribute deliberately does not auto-compute ---------------- + +_MUTATES = ( + "in-place mutator; auto-computing would mutate a throwaway temporary, " + "not the caller's object" +) +_CONSTRUCTS = ( + "constructor (classmethod) that builds a new object from external data; " + "not an accessor on a computed result" +) +_MATERIALIZES = ( + "produces a new concrete object; an expression is materialized with .new(), " + "so an auto-compute property would be redundant" +) +_METADATA = ( + "cheap metadata known without materializing; exposed natively on expressions " + "via BaseExpression, so it must not force a compute" +) +_SCALAR_STORAGE = ( + "storage-backing flag (C scalar vs GrB_Scalar); ScalarExpression and " + "ScalarIndexExpr define it natively since they know how their result will " + "materialize, so it must not force a compute" +) +_RAISES_MATERIALIZE = ( + "raises to require explicit materialization of the lazy transposed view " + "(np.asarray / bool); mirrors Matrix-expression behavior" +) + +# Every public name of a concrete type that intentionally is NOT auto-computed. +# Adding a public method almost never belongs here; it belongs in the automethods +# name sets. This table is for the deliberate exceptions only. +OPT_OUT = { + "Scalar": { + "clear": _MUTATES, + "update": _MUTATES, + "dup": _MATERIALIZES, + "from_value": _CONSTRUCTS, + "dtype": _METADATA, + "ndim": _METADATA, + "shape": _METADATA, + "is_cscalar": _SCALAR_STORAGE, + "is_grbscalar": _SCALAR_STORAGE, + }, + "Vector": { + "build": _MUTATES, + "clear": _MUTATES, + "resize": _MUTATES, + "update": _MUTATES, + "dup": _MATERIALIZES, + "from_coo": _CONSTRUCTS, + "from_dense": _CONSTRUCTS, + "from_dict": _CONSTRUCTS, + "from_pairs": _CONSTRUCTS, + "from_scalar": _CONSTRUCTS, + "dtype": _METADATA, + "ndim": _METADATA, + "shape": _METADATA, + "size": _METADATA, + }, + "Matrix": { + "build": _MUTATES, + "clear": _MUTATES, + "resize": _MUTATES, + "setdiag": _MUTATES, + "update": _MUTATES, + "dup": _MATERIALIZES, + "from_coo": _CONSTRUCTS, + "from_csc": _CONSTRUCTS, + "from_csr": _CONSTRUCTS, + "from_dcsc": _CONSTRUCTS, + "from_dcsr": _CONSTRUCTS, + "from_dense": _CONSTRUCTS, + "from_dicts": _CONSTRUCTS, + "from_edgelist": _CONSTRUCTS, + "from_scalar": _CONSTRUCTS, + "dtype": _METADATA, + "ncols": _METADATA, + "ndim": _METADATA, + "nrows": _METADATA, + "shape": _METADATA, + }, + "TransposedMatrix": { + "dup": _MATERIALIZES, + "new": _MATERIALIZES, + "dtype": _METADATA, + "ncols": _METADATA, + "ndim": _METADATA, + "nrows": _METADATA, + "shape": _METADATA, + "__array__": _RAISES_MATERIALIZE, + "__bool__": _RAISES_MATERIALIZE, + }, +} + +# Concrete type + the expression classes the generator targets for it: the plain +# expression, the index expression, and the infix expression. Coverage is the +# intersection across all three, so a name emitted onto some but not the infix +# class is treated as uncovered. +# TransposedMatrix has no expression class of its own; as a read-only Matrix view +# it shares Matrix's generated surface. +_REGISTRY = { + "Scalar": (Scalar, (ScalarExpression, ScalarIndexExpr, ScalarInfixExpr)), + "Vector": (Vector, (VectorExpression, VectorIndexExpr, VectorInfixExpr)), + "Matrix": (Matrix, (MatrixExpression, MatrixIndexExpr, MatrixInfixExpr)), + "TransposedMatrix": (TransposedMatrix, (MatrixExpression, MatrixIndexExpr, MatrixInfixExpr)), +} + + +# --- dunder scope helpers --------------------------------------------------- + + +class _Baseline: + __slots__ = () + + +# Object/interpreter machinery dunders. Computed from a trivial class so that +# version-specific additions (e.g. __firstlineno__, __static_attributes__) are +# absorbed automatically rather than hard-coded. +_MACHINERY_DUNDERS = ( + set(dir(_Baseline)) + | set(vars(_Baseline)) + | { + "__del__", + "__weakref__", + "__dict__", + "__reduce__", + "__reduce_ex__", + "__getstate__", + "__setstate__", + # copyreg._slotnames() caches this on the class the first time an + # instance is pickled, so its presence depends on whether test_pickle + # ran first in the randomized test order. + "__slotnames__", + "__networkx_backend__", + "__networkx_plugin__", + } +) + +# Arithmetic / comparison / item-mutation sugar. These build lazy infix +# expressions (or mutate in place) and are handled by infixmethods.py, not the +# value-forwarding automethods path, so they are out of scope for this sweep. +_OPERATOR_SUGAR_DUNDERS = { + "__add__", + "__radd__", + "__sub__", + "__rsub__", + "__mul__", + "__rmul__", + "__truediv__", + "__rtruediv__", + "__floordiv__", + "__rfloordiv__", + "__mod__", + "__rmod__", + "__pow__", + "__rpow__", + "__divmod__", + "__rdivmod__", + "__xor__", + "__rxor__", + "__neg__", + "__abs__", + "__invert__", + "__lt__", + "__le__", + "__gt__", + "__ge__", + "__setitem__", + "__delitem__", +} + + +# --- introspection ---------------------------------------------------------- + + +def _accessible(cls, name): + """True if ``cls.name`` resolves without raising. + + ``dir()`` lists names that are not usable (e.g. ``ss`` under the + suitesparse-vanilla backend raises on access); those are not part of the + public surface a caller can rely on. + """ + try: + getattr(cls, name) + except Exception: + # Any failure to access means the name is not a usable public attribute. + return False + return True + + +def _public_surface(cls): + return {n for n in dir(cls) if not n.startswith("_") and _accessible(cls, n)} + + +def _own_dunders(cls): + return {n for n in vars(cls) if n.startswith("__") and n.endswith("__")} + + +def _generated_coverage_one(expr_class): + """Names the generator emitted onto a single ``expr_class``. + + A name counts as covered when the class exposes it as a property whose getter + is defined in the automethods module, or as a callable copied from that module + (the ``__iadd__``-style guards). This reads the actual generated surface, so it + tracks the name sets without duplicating them. + """ + names = set() + for klass in expr_class.__mro__: + for name, value in vars(klass).items(): + func = value.fget if isinstance(value, property) else value + if callable(func) and getattr(func, "__module__", None) == _AUTOMETHODS_MODULE: + names.add(name) + return names + + +def _generated_coverage(*expr_classes): + """Names the generator emitted onto EVERY class in ``expr_classes``. + + The intersection: a name is covered only when it is reachable on all variants + (plain expression, index expression, and infix expression). One emitted onto + some but not others counts as uncovered, which is how a name missing from just + the infix surface is caught. + """ + per_class = [_generated_coverage_one(cls) for cls in expr_classes] + return set.intersection(*per_class) if per_class else set() + + +def _coverage(label): + _concrete, expr_classes = _REGISTRY[label] + return _generated_coverage(*expr_classes) + + +def _fix_hint(label, names): + _concrete, expr_classes = _REGISTRY[label] + expr_names = " / ".join(c.__name__ for c in dict.fromkeys(expr_classes)) + return ( + f"{label} has public attribute(s) not reachable on its expression " + f"classes ({expr_names}) and not listed in OPT_OUT[{label!r}]:\n" + f" {sorted(names)}\n\n" + "If these should auto-compute on expressions, add each name to the " + "matching set in graphblas/core/automethods.py (the sets near line 347) " + "and regenerate with:\n" + " python scripts/autogenerate.py\n\n" + "If they must NOT auto-compute (an in-place mutator, a constructor, or " + f"cheap metadata), add them to OPT_OUT[{label!r}] in this file with a " + "one-line reason." + ) + + +# --- tests ------------------------------------------------------------------ + + +@pytest.mark.parametrize("label", list(_REGISTRY)) +def test_public_methods_covered_or_opted_out(label): + concrete, _expr_classes = _REGISTRY[label] + coverage = _coverage(label) + opt_out = OPT_OUT[label] + uncovered = _public_surface(concrete) - coverage - set(opt_out) + assert not uncovered, _fix_hint(label, uncovered) + + +@pytest.mark.parametrize("label", list(_REGISTRY)) +def test_relevant_dunders_covered_or_opted_out(label): + concrete, _expr_classes = _REGISTRY[label] + coverage = _coverage(label) + opt_out = OPT_OUT[label] + candidates = ( + _own_dunders(concrete) + - _MACHINERY_DUNDERS + - _OPERATOR_SUGAR_DUNDERS + - coverage + - set(opt_out) + ) + assert not candidates, _fix_hint(label, candidates) + + +@pytest.mark.parametrize("label", ["Scalar", "Vector", "Matrix"]) +def test_no_stale_generated_names(label): + concrete, _expr_classes = _REGISTRY[label] + # _get_value is the auto-compute helper itself, not a mirror of a concrete + # attribute, so it is expected not to exist on the concrete type. + stale = {n for n in _coverage(label) if n != "_get_value" and not hasattr(concrete, n)} + assert not stale, ( + f"Generated name(s) on the {label} expression classes no longer exist " + f"on {label}: {sorted(stale)}. A concrete method was renamed or removed " + "without updating the sets in graphblas/core/automethods.py; update the " + "sets and rerun `python scripts/autogenerate.py`." + ) + + +@pytest.mark.parametrize("label", list(_REGISTRY)) +def test_opt_out_entries_are_live(label): + concrete, _expr_classes = _REGISTRY[label] + opt_out = OPT_OUT[label] + surface = _public_surface(concrete) | _own_dunders(concrete) + missing = {n for n in opt_out if n not in surface} + assert not missing, ( + f"OPT_OUT[{label!r}] lists name(s) that are not attributes of {label}: " + f"{sorted(missing)}. Remove the stale entries (the method was renamed or " + "removed)." + ) + coverage = _coverage(label) + redundant = {n for n in opt_out if n in coverage} + assert not redundant, ( + f"OPT_OUT[{label!r}] lists name(s) that ARE auto-computed and so need no " + f"opt-out: {sorted(redundant)}. Remove them from OPT_OUT." + ) + + +@pytest.mark.parametrize("label", list(_REGISTRY)) +def test_opt_out_reasons_present(label): + empty = {n for n, reason in OPT_OUT[label].items() if not reason or not reason.strip()} + assert not empty, f"OPT_OUT[{label!r}] entries need a non-empty reason: {sorted(empty)}" diff --git a/graphblas/tests/test_core.py b/graphblas/tests/test_core.py index a15da6f54..859c5c21c 100644 --- a/graphblas/tests/test_core.py +++ b/graphblas/tests/test_core.py @@ -94,3 +94,48 @@ def test_packages(): def test_index_max(): assert gb.MAX_SIZE == 2**60 # True for all current backends + + +def test_config_attribute_assignment_raises(): + # Setting an option by plain attribute assignment used to silently no-op: + # donfig has no __setattr__ for options, so it created a dead instance + # attribute and left the real option alone. It should raise and name the + # canonical API instead. + orig = gb.config["autocompute"] + try: + with pytest.raises(AttributeError, match=r"config\.set\(autocompute"): + gb.config.autocompute = not orig + assert gb.config["autocompute"] == orig # the failed write changed nothing + # An unknown name is a mistaken write too, not a new attribute + with pytest.raises(AttributeError, match="not_a_real_option"): + gb.config.not_a_real_option = 5 + assert "not_a_real_option" not in vars(gb.config) + # Canonical idioms still work: read by item, write by set() + gb.config.set(autocompute=not orig) + assert gb.config["autocompute"] == (not orig) + with gb.config.set(autocompute=orig): + assert gb.config["autocompute"] == orig + assert gb.config["autocompute"] == (not orig) # restored on exit + finally: + gb.config.set(autocompute=orig) + + +def test_config_donfig_attrs_are_derived(monkeypatch): + # The allowlist of writable attributes must come from what donfig's + # __init__ actually set, not a hardcoded list. Simulate a future donfig + # that grows a field: the derived allowlist absorbs it, where a + # hardcoded list would start rejecting donfig's own writes. + import donfig + + real_init = donfig.Config.__init__ + + def init_with_extra_field(self, *args, **kwargs): + real_init(self, *args, **kwargs) + self.hypothetical_future_field = 1 + + monkeypatch.setattr(donfig.Config, "__init__", init_with_extra_field) + probe = type(gb.config)("test_derived_allowlist") + assert "hypothetical_future_field" in probe._donfig_attrs + probe.hypothetical_future_field = 2 # writable, not rejected + with pytest.raises(AttributeError, match="autocompute"): + probe.autocompute = False diff --git a/graphblas/tests/test_fastpath_parity.py b/graphblas/tests/test_fastpath_parity.py new file mode 100644 index 000000000..e34fb6ee0 --- /dev/null +++ b/graphblas/tests/test_fastpath_parity.py @@ -0,0 +1,961 @@ +"""Parity tests for the scalar-access fast paths. + +Several hot scalar-access operations grew fast paths that bypass the general +expression machinery: ``Vector.get`` / ``Matrix.get``, ``index in obj`` +(``__contains__``), integer-key ``obj[i] = value`` (``__setitem__``), +``obj[i].new()`` on a scalar index expression, and the plain-int lane in +``parse_index``. Each fast path is only correct if it produces exactly what the +slower general path would. These tests pin that equivalence: for every case we +run the shipped fast path and an inlined copy of the pre-fast-path reference, +then require the results (values, dtypes, and raised exceptions with their +messages) to match. + +The references here are deliberately self-contained rather than imported from +the throwaway benchmark scripts they were distilled from, so a refactor of the +production code can never silently drag the reference along with it. +""" + +import numpy as np +import pytest + +import graphblas as gb +from graphblas import Matrix, Vector, binary, dtypes +from graphblas.core import _supports_udfs as supports_udfs # noqa: F401 +from graphblas.core.expr import Updater +from graphblas.core.recorder import Recorder +from graphblas.core.scalar import Scalar + +# Builtin dtypes to sweep. Complex is SuiteSparse-only, so gate it the same way +# the dtype tests do; on the vanilla backend the parametrization simply omits it. +BUILTIN_DTYPES = [ + dtypes.BOOL, + dtypes.INT8, + dtypes.INT16, + dtypes.INT32, + dtypes.INT64, + dtypes.UINT8, + dtypes.UINT16, + dtypes.UINT32, + dtypes.UINT64, + dtypes.FP32, + dtypes.FP64, +] +if dtypes._supports_complex: + BUILTIN_DTYPES += [dtypes.FC32, dtypes.FC64] + +DTYPE_IDS = [dt.name for dt in BUILTIN_DTYPES] + + +def _hit_value(dt): + """A value that fits every builtin dtype (matches the source harnesses).""" + if "FC" in dt.name: + return 3 + 4j + if dt == dtypes.BOOL: + return True + return 3 + + +def _capture(fn): + """Return ``("val", result)`` or ``("exc", type_name, message)``.""" + try: + return ("val", fn()) + except Exception as e: + return ("exc", type(e).__name__, str(e)) + + +class _IndexLike: + """Implements ``__index__`` without being an int; must not take any fast lane. + + ``parse_index`` rejects it, so a fast lane that duck-types ``__index__`` + would diverge from the expression path (gh: scalar fast-path parity). + """ + + def __init__(self, value): + self.value = value + + def __index__(self): + return self.value + + def __repr__(self): + return f"_IndexLike({self.value})" + + +# --------------------------------------------------------------------------- +# Vector.get / Matrix.get +# --------------------------------------------------------------------------- +def _get_vector_slow(v, index, default=None): + expr = v[index] + if not expr._is_scalar: + raise ValueError("Bad index in Vector.get(...)") + rv = expr.new().value + return default if rv is None else rv + + +def _get_matrix_slow(A, r, c, default=None): + rv = A[r, c].new().value + return default if rv is None else rv + + +def _assert_get_equal(got, expected): + msg = f"got {got!r} ({type(got).__name__}), expected {expected!r} ({type(expected).__name__})" + assert got == expected, msg + assert type(got) is type(expected), msg + + +@pytest.mark.parametrize("dt", BUILTIN_DTYPES, ids=DTYPE_IDS) +def test_get_dtype_parity(dt): + val = _hit_value(dt) + v = Vector(dt, 10) + v[2] = val + _assert_get_equal(v.get(2), _get_vector_slow(v, 2)) + _assert_get_equal(v.get(3), _get_vector_slow(v, 3)) + _assert_get_equal(v.get(3, -1), _get_vector_slow(v, 3, -1)) + A = Matrix(dt, 5, 7) + A[1, 2] = val + _assert_get_equal(A.get(1, 2), _get_matrix_slow(A, 1, 2)) + _assert_get_equal(A.get(0, 0, 99), _get_matrix_slow(A, 0, 0, 99)) + + +def test_get_index_types(): + v = Vector.from_coo([0, 5, 9], [1.5, 2.5, 3.5], size=10) + _assert_get_equal(v.get(np.int64(5)), 2.5) + _assert_get_equal(v.get(np.uint8(5)), 2.5) + _assert_get_equal(v.get(-5), 2.5) + _assert_get_equal(v.get(-2, "d"), "d") + # empty-scalar semantics: a miss with no default is None + assert v.get(1) is None + + +@pytest.mark.parametrize("idx", [10, -11, 1000]) +def test_get_out_of_range(idx): + """Out-of-range raises IndexError with the same message on both paths.""" + v = Vector.from_coo([0, 5, 9], [1.5, 2.5, 3.5], size=10) + with pytest.raises(IndexError) as fast: + v.get(idx) + with pytest.raises(IndexError) as slow: + _get_vector_slow(v, idx) + assert str(fast.value) == str(slow.value) + + +@pytest.mark.parametrize( + "bad", [1.5, "x", None, [1, 2], slice(None), 2.0, True, np.array(3), _IndexLike(3)] +) +def test_get_non_integer_index(bad): + """Non-integer indices raise the same exception type as the reference path.""" + v = Vector.from_coo([0, 5, 9], [1.5, 2.5, 3.5], size=10) + fast = _capture(lambda: v.get(bad)) + slow = _capture(lambda: _get_vector_slow(v, bad)) + assert fast[:2] == slow[:2], f"fast={fast} slow={slow}" + + +@pytest.mark.parametrize( + ("r", "c"), + [ + (np.array(1), 2), + (1, np.array(2)), + (np.array(1), np.array(2)), + (_IndexLike(1), 2), + (1, _IndexLike(2)), + ], +) +def test_get_matrix_non_integer_index(r, c): + """Index-like objects that parse_index rejects must fail identically in get.""" + A = Matrix.from_coo([0, 1, 4], [1, 2, 6], [10, 20, 30], nrows=5, ncols=7) + fast = _capture(lambda: A.get(r, c)) + slow = _capture(lambda: _get_matrix_slow(A, r, c)) + assert fast == slow, f"fast={fast} slow={slow}" + + +def test_get_matrix_negative_and_range(): + A = Matrix.from_coo([0, 1, 4], [1, 2, 6], [10, 20, 30], nrows=5, ncols=7) + _assert_get_equal(A.get(-1, -1, "d"), 30) + _assert_get_equal(A.get(-1, 6), 30) + _assert_get_equal(A.get(-2, -2, "d"), "d") + with pytest.raises(IndexError): + A.get(5, 0) + with pytest.raises(IndexError): + A.get(0, 7) + + +def test_get_transposed_matrix(): + """TransposedMatrix reuses Matrix.get and must extract the mirrored element.""" + # Non-square and non-symmetric, so a row/column swap cannot go unnoticed. + A = Matrix.from_coo([0, 0, 1, 2], [1, 3, 0, 2], [10, 30, 20, 40], nrows=3, ncols=4) + AT = A.T + assert AT.shape == (4, 3) + for r in range(4): + for c in range(3): + _assert_get_equal(AT.get(r, c, "d"), _get_matrix_slow(AT, r, c, "d")) + # The mirrored element is the same one A sees with the axes swapped. + _assert_get_equal(AT.get(r, c, "d"), A.get(c, r, "d")) + # Out-of-range is judged against the transposed dimensions. + with pytest.raises(IndexError): + AT.get(4, 0) + with pytest.raises(IndexError): + AT.get(0, 3) + # The setitem fast path is unreachable through a transposed view. + assert not hasattr(type(AT), "__setitem__") + + +def test_get_udt_fallback(): + udt = gb.dtypes.register_anonymous( + np.dtype([("fx", np.int64), ("fy", np.float64)]), "GetFastPathProbe" + ) + u = Vector(udt, 3) + u[1] = (7, 2.5) + r = u.get(1) + assert r["fx"] == 7, r + assert r["fy"] == 2.5, r + assert u.get(0, "dflt") == "dflt" + + +def test_get_recorder_fallback(): + """An active Recorder must see the extract call, i.e. the fast path defers.""" + v = Vector.from_coo([0], [1.5], size=4) + with Recorder() as rec: + assert v.get(0) == 1.5 + data = "".join(rec.data) + assert "extractElement" in data, f"recorder missed call: {data!r}" + + +def test_get_pending_value(): + """A value set in non-blocking mode is visible before an explicit wait.""" + v = Vector(float, 100) + v[3] = 2.25 # pending setElement in non-blocking mode + _assert_get_equal(v.get(3), 2.25) + + +# --------------------------------------------------------------------------- +# __contains__ (index in obj) +# --------------------------------------------------------------------------- +def _contains_vector_slow(v, index): + extractor = v[index] + if not extractor._is_scalar: + raise TypeError( + f"Invalid index to Vector contains: {index!r}. An integer is expected. " + "Doing `index in my_vector` checks whether a value is present at that index." + ) + scalar = extractor.new(name="s_contains") + return not scalar._is_empty + + +def _contains_matrix_slow(A, index): + extractor = A[index] + if not extractor._is_scalar: + raise TypeError( + f"Invalid index to Matrix contains: {index!r}. A 2-tuple of ints is expected. " + "Doing `(i, j) in my_matrix` checks whether a value is present at that index." + ) + scalar = extractor.new(name="s_contains") + return not scalar._is_empty + + +@pytest.mark.parametrize("dt", BUILTIN_DTYPES, ids=DTYPE_IDS) +def test_contains_dtype_parity(dt): + val = _hit_value(dt) + v = Vector(dt, 10) + v[2] = val + assert (2 in v) == _contains_vector_slow(v, 2) + assert (3 in v) == _contains_vector_slow(v, 3) + # a zero / false value is still "present" + v[4] = False if dt == dtypes.BOOL else 0 + assert (4 in v) == _contains_vector_slow(v, 4) + A = Matrix(dt, 5, 7) + A[1, 2] = val + assert ((1, 2) in A) == _contains_matrix_slow(A, (1, 2)) + assert ((0, 0) in A) == _contains_matrix_slow(A, (0, 0)) + + +def test_contains_index_types(): + v = Vector.from_coo([0, 5, 9], [1.5, 2.5, 3.5], size=10) + assert (np.int64(5) in v) == _contains_vector_slow(v, np.int64(5)) + assert (np.uint8(5) in v) == _contains_vector_slow(v, np.uint8(5)) + assert (np.int64(6) in v) == _contains_vector_slow(v, np.int64(6)) + assert (-5 in v) == _contains_vector_slow(v, -5) + assert (-2 in v) == _contains_vector_slow(v, -2) + assert (-1 in v) == _contains_vector_slow(v, -1) + + +@pytest.mark.parametrize("idx", [10, -11, 1000, -1000]) +def test_contains_out_of_range(idx): + """Out-of-range raises the identical IndexError through either lane.""" + v = Vector.from_coo([0, 5, 9], [1.5, 2.5, 3.5], size=10) + fast = _capture(lambda: idx in v) + slow = _capture(lambda: _contains_vector_slow(v, idx)) + assert fast == slow, f"fast={fast} slow={slow}" + assert fast[:2] == ("exc", "IndexError"), fast + + +@pytest.mark.parametrize("b", [True, False]) +def test_contains_bool_index(b): + v = Vector.from_coo([0, 5, 9], [1.5, 2.5, 3.5], size=10) + fast = _capture(lambda: b in v) + slow = _capture(lambda: _contains_vector_slow(v, b)) + assert fast == slow, f"fast={fast} slow={slow}" + + +@pytest.mark.parametrize( + "bad", [1.5, "x", None, [1, 2], slice(None), 2.0, (1, 2), np.array(3), _IndexLike(3)] +) +def test_contains_non_integer_index(bad): + """Non-integer indices raise the same exception type AND message.""" + v = Vector.from_coo([0, 5, 9], [1.5, 2.5, 3.5], size=10) + fast = _capture(lambda: bad in v) + slow = _capture(lambda: _contains_vector_slow(v, bad)) + assert fast == slow, f"fast={fast} slow={slow}" + + +def test_contains_matrix_negative_and_range(): + A = Matrix.from_coo([0, 1, 4], [1, 2, 6], [10, 20, 30], nrows=5, ncols=7) + assert ((-1, -1) in A) == _contains_matrix_slow(A, (-1, -1)) + assert ((-1, 6) in A) == _contains_matrix_slow(A, (-1, 6)) + assert ((-2, -2) in A) == _contains_matrix_slow(A, (-2, -2)) + for idx in [(5, 0), (0, 7), (-6, 0), (0, -8), (100, 100)]: + fast = _capture(lambda idx=idx: idx in A) + slow = _capture(lambda idx=idx: _contains_matrix_slow(A, idx)) + assert fast == slow, f"{idx}: fast={fast} slow={slow}" + assert fast[:2] == ("exc", "IndexError"), (idx, fast) + + +def test_contains_transposed_matrix(): + """TransposedMatrix shares Matrix.__contains__; mirrored indices must agree.""" + A = Matrix.from_coo([0, 1, 4], [1, 2, 6], [10, 20, 30], nrows=5, ncols=7) + AT = A.T + for idx in [(0, 0), (6, 0), (2, 1), (6, 4), (-1, -1)]: + assert (idx in AT) == _contains_matrix_slow(AT, idx), idx + fast = _capture(lambda: (0, 6) in AT) + slow = _capture(lambda: _contains_matrix_slow(AT, (0, 6))) + assert fast == slow, f"fast={fast} slow={slow}" + assert fast[:2] == ("exc", "IndexError"), fast + + +@pytest.mark.parametrize( + "bad", + [ + 5, + (1, 2, 3), + (1.5, 2), + ("a", "b"), + np.int64(3), + (np.array(1), np.array(2)), + (1, np.array(2)), + (_IndexLike(1), 2), + ], +) +def test_contains_matrix_bad_index(bad): + A = Matrix.from_coo([0, 1, 4], [1, 2, 6], [10, 20, 30], nrows=5, ncols=7) + fast = _capture(lambda: bad in A) + slow = _capture(lambda: _contains_matrix_slow(A, bad)) + assert fast == slow, f"fast={fast} slow={slow}" + + +@pytest.mark.parametrize("pair", [(True, 0), (0, True), (True, True)]) +def test_contains_matrix_bool_pair(pair): + A = Matrix.from_coo([0, 1, 4], [1, 2, 6], [10, 20, 30], nrows=5, ncols=7) + fast = _capture(lambda: pair in A) + slow = _capture(lambda: _contains_matrix_slow(A, pair)) + assert fast == slow, f"fast={fast} slow={slow}" + + +def test_contains_udt_fallback(): + udt = gb.dtypes.register_anonymous( + np.dtype([("cx", np.int64), ("cy", np.float64)]), "ContainsFastPathProbe" + ) + u = Vector(udt, 3) + u[1] = (7, 2.5) + assert (1 in u) == _contains_vector_slow(u, 1) + assert (0 in u) == _contains_vector_slow(u, 0) + fast = _capture(lambda: 5 in u) + slow = _capture(lambda: _contains_vector_slow(u, 5)) + assert fast == slow, f"fast={fast} slow={slow}" + assert fast[:2] == ("exc", "IndexError"), fast + Mu = Matrix(udt, 3, 3) + Mu[1, 1] = (7, 2.5) + assert ((1, 1) in Mu) == _contains_matrix_slow(Mu, (1, 1)) + assert ((0, 0) in Mu) == _contains_matrix_slow(Mu, (0, 0)) + + +def test_contains_recorder_fallback(): + v = Vector.from_coo([0], [1.5], size=4) + with Recorder() as rec: + assert 0 in v + assert "extractElement" in "".join(rec.data) + A = Matrix.from_coo([0], [0], [1.5], nrows=4, ncols=4) + with Recorder() as rec: + assert (0, 0) in A + assert "extractElement" in "".join(rec.data) + + +def test_contains_pending_value(): + v = Vector(float, 100) + v[3] = 2.25 + assert (3 in v) is True + assert (4 in v) is False + + +# --------------------------------------------------------------------------- +# __setitem__ with an integer key +# --------------------------------------------------------------------------- +def _set_vector_slow(v, key, val): + Updater(v, opts={})[key] = val + + +def _set_matrix_slow(A, key, val): + Updater(A, opts={})[key] = val + + +def _assert_setitem_vector(dt, key, val, size=10): + a = Vector(dt, size) + b = Vector(dt, size) + ea = _capture(lambda: a.__setitem__(key, val)) + eb = _capture(lambda: _set_vector_slow(b, key, val)) + # On success both capture ("val", None); on error both capture the full + # ("exc", type, message), so equality pins exception type AND message. + assert ea == eb, f"exc fast={ea} slow={eb}" + if ea[0] == "val": + assert a.isequal(b, check_dtype=True), f"value fast={a.to_coo()} slow={b.to_coo()}" + + +def _assert_setitem_matrix(dt, key, val, nrows=5, ncols=7): + a = Matrix(dt, nrows, ncols) + b = Matrix(dt, nrows, ncols) + ea = _capture(lambda: a.__setitem__(key, val)) + eb = _capture(lambda: _set_matrix_slow(b, key, val)) + assert ea == eb, f"exc fast={ea} slow={eb}" + if ea[0] == "val": + assert a.isequal(b, check_dtype=True), f"value fast={a.to_coo()} slow={b.to_coo()}" + + +def _setitem_natural_value(dt): + if "FC" in dt.name: + return 3 + 4j + if dt == dtypes.BOOL: + return True + if "FP" in dt.name: + return 3.5 + return 3 + + +# Cross-type coercion: the value's Python type differs from the container dtype. +# The point is that whatever SuiteSparse does on the cast, the fast path does it +# too. Complex-target cases are gated on complex support. +CROSS_COERCION = [ + (dtypes.FP64, 3), # int -> FP64 + (dtypes.FP32, 3), # int -> FP32 + (dtypes.INT64, 3.9), # float -> INT64 (truncation via SS cast) + (dtypes.INT32, -2.5), # float -> INT32 + (dtypes.INT8, True), # bool -> INT8 + (dtypes.UINT8, True), # bool -> UINT8 + (dtypes.FP64, True), # bool -> FP64 + (dtypes.UINT8, -1), # negative int -> unsigned (wrap per SS) + (dtypes.UINT8, 256), # int overflow of the container (SS casts) + (dtypes.INT8, 200), # int overflow of the container + (dtypes.INT64, 2**62), # large int that fits int64 + (dtypes.INT64, 2**63), # overflow int64 -> OverflowError on both + (dtypes.UINT64, 2**63), # fits uint64 + (dtypes.FP64, 1e308), # large float + (dtypes.FP32, 1e308), # overflow FP32 -> inf via SS cast +] +if dtypes._supports_complex: + CROSS_COERCION += [ + (dtypes.FP64, 3 + 0j), # complex -> real (whatever SS does) + (dtypes.FC64, 3), # int -> FC64 + (dtypes.FC64, 3.5), # float -> FC64 + ] + +CROSS_IDS = [f"{dt.name}<-{val!r}" for dt, val in CROSS_COERCION] + + +@pytest.mark.parametrize("dt", BUILTIN_DTYPES, ids=DTYPE_IDS) +def test_setitem_dtype_natural(dt): + val = _setitem_natural_value(dt) + _assert_setitem_vector(dt, 2, val) + _assert_setitem_matrix(dt, (1, 2), val) + + +@pytest.mark.parametrize(("dt", "val"), CROSS_COERCION, ids=CROSS_IDS) +def test_setitem_cross_coercion(dt, val): + _assert_setitem_vector(dt, 4, val) + _assert_setitem_matrix(dt, (2, 3), val) + + +@pytest.mark.parametrize("dt", [dtypes.FP64, dtypes.INT64, dtypes.BOOL], ids=lambda dt: dt.name) +def test_setitem_overwrite(dt): + a = Vector(dt, 6) + b = Vector(dt, 6) + a[1] = 1 + _set_vector_slow(b, 1, 1) + a[1] = 5 # overwrite + _set_vector_slow(b, 1, 5) + assert a.isequal(b, check_dtype=True) + + +def test_setitem_numpy_index(): + for key in [np.int64(3), np.uint8(3), np.int32(3)]: + _assert_setitem_vector(dtypes.FP64, key, 2.5) + _assert_setitem_matrix(dtypes.FP64, (np.int64(1), np.int32(2)), 2.5) + + +def test_setitem_negative_index(): + _assert_setitem_vector(dtypes.FP64, -1, 9.0) + _assert_setitem_vector(dtypes.FP64, -10, 9.0) + _assert_setitem_matrix(dtypes.FP64, (-1, -1), 9.0) + _assert_setitem_matrix(dtypes.FP64, (-2, 3), 9.0) + + +@pytest.mark.parametrize("key", [10, -11, 1000, -1000]) +def test_setitem_vector_out_of_range(key): + _assert_setitem_vector(dtypes.FP64, key, 1.0) + + +@pytest.mark.parametrize("key", [(5, 0), (0, 7), (-6, 0), (0, -8), (100, 3), (2, 100)]) +def test_setitem_matrix_out_of_range(key): + _assert_setitem_matrix(dtypes.FP64, key, 1.0) + + +@pytest.mark.parametrize( + "key", [1.5, "x", None, slice(None), [1, 2], (1, 2), np.array(3), _IndexLike(3)] +) +def test_setitem_vector_bad_key(key): + _assert_setitem_vector(dtypes.FP64, key, 1.0) + + +@pytest.mark.parametrize( + "key", + [ + 5, + (1, 2, 3), + (1.5, 2), + ("a", "b"), + 1.5, + slice(None), + (np.array(1), np.array(2)), + (1, np.array(2)), + (_IndexLike(1), 2), + ], +) +def test_setitem_matrix_bad_key(key): + _assert_setitem_matrix(dtypes.FP64, key, 1.0) + + +def test_setitem_bool_key(): + # bool is treated as an int index by parse_index; the fast path falls back. + _assert_setitem_vector(dtypes.FP64, True, 1.0) + _assert_setitem_vector(dtypes.FP64, False, 1.0) + _assert_setitem_matrix(dtypes.FP64, (True, 0), 1.0) + + +def test_setitem_value_fallbacks(): + # numpy scalar, Scalar, None, and str all route through the slow path. + _assert_setitem_vector(dtypes.FP64, 2, np.int32(7)) + _assert_setitem_vector(dtypes.INT64, 2, np.float64(7.9)) + _assert_setitem_vector(dtypes.FP64, 2, Scalar.from_value(7.5)) + _assert_setitem_vector(dtypes.FP64, 2, None) + _assert_setitem_vector(dtypes.FP64, 2, "bad") + _assert_setitem_matrix(dtypes.FP64, (1, 1), np.int32(7)) + _assert_setitem_matrix(dtypes.FP64, (1, 1), None) + + +# Both halves compare UDT vectors with ``isequal``, which needs ``binary.eq`` +# compiled for the UDT. The sibling UDT tests here never call ``isequal``, so +# they keep running on builds without numba. +@pytest.mark.skipif("not supports_udfs") +def test_setitem_udt_fallback(): + udt = gb.dtypes.register_anonymous( + np.dtype([("sx", np.int64), ("sy", np.float64)]), "SetitemFastPathProbe" + ) + _assert_setitem_vector(udt, 1, None) # None clears; both fall back + au = Vector(udt, 3) + bu = Vector(udt, 3) + au[1] = (7, 2.5) + _set_vector_slow(bu, 1, (7, 2.5)) + assert au.isequal(bu, check_dtype=True) + + +def test_setitem_recorder_fallback(): + v = Vector(dtypes.FP64, 4) + with Recorder() as rec: + v[0] = 1.5 + assert "setElement" in "".join(rec.data) + A = Matrix(dtypes.FP64, 4, 4) + with Recorder() as rec: + A[0, 0] = 1.5 + assert "setElement" in "".join(rec.data) + + +def test_setitem_update_forms(): + """``obj[i] << x`` and masked / accum forms stay on the Updater path.""" + w = Vector(dtypes.FP64, 5) + w[2] << 3.5 + assert w.get(2) == 3.5 + w(accum=binary.plus)[2] << 1.5 + assert w.get(2) == 5.0 + A = Matrix(dtypes.FP64, 4, 4) + A[1, 1] << 3.5 + assert A.get(1, 1) == 3.5 + m = Vector(dtypes.FP64, 5) + m[:] = 1.0 + mask = Vector(dtypes.BOOL, 5) + mask[0] = True + mask[2] = True + m(mask.V)[:] = 9.0 + assert m.get(0) == 9.0 + assert m.get(1) == 1.0 + assert m.get(2) == 9.0 + + +def test_setitem_empty_vector(): + e = Vector(dtypes.INT32, 3) + e[1] = 42 + assert e.get(1) == 42 + assert e.nvals == 1 + + +# --------------------------------------------------------------------------- +# ScalarIndexExpr.new() +# --------------------------------------------------------------------------- +def _new_slow(expr, dtype=None, is_cscalar=None, name=None, **opts): + if is_cscalar is None: + is_cscalar = False + return expr.parent._extract_element( + expr.resolved_indexes, dtype, opts, is_cscalar=is_cscalar, name=name + ) + + +def _assert_scalars_equal(fast, slow): + assert fast.dtype == slow.dtype, f"dtype {fast.dtype} vs {slow.dtype}" + assert fast.is_cscalar == slow.is_cscalar, f"is_cscalar {fast.is_cscalar} vs {slow.is_cscalar}" + assert fast.is_empty == slow.is_empty, f"is_empty {fast.is_empty} vs {slow.is_empty}" + fv, sv = fast.value, slow.value + assert fv == sv or (fv is None and sv is None), f"value {fv!r} vs {sv!r}" + assert type(fv) is type(sv), f"value type {type(fv).__name__} vs {type(sv).__name__}" + + +@pytest.mark.parametrize("dt", BUILTIN_DTYPES, ids=DTYPE_IDS) +def test_scalarnew_dtype_parity(dt): + val = _hit_value(dt) + v = Vector(dt, 10) + v[2] = val + _assert_scalars_equal(v[2].new(), _new_slow(v[2])) + _assert_scalars_equal(v[3].new(), _new_slow(v[3])) + A = Matrix(dt, 5, 7) + A[1, 2] = val + _assert_scalars_equal(A[1, 2].new(), _new_slow(A[1, 2])) + _assert_scalars_equal(A[0, 0].new(), _new_slow(A[0, 0])) + + +def test_scalarnew_negative_and_numpy(): + v = Vector.from_coo([0, 5, 9], [1.5, 2.5, 3.5], size=10) + _assert_scalars_equal(v[-1].new(), _new_slow(v[-1])) + _assert_scalars_equal(v[-5].new(), _new_slow(v[-5])) + _assert_scalars_equal(v[-2].new(), _new_slow(v[-2])) + _assert_scalars_equal(v[np.int64(5)].new(), _new_slow(v[np.int64(5)])) + + +def test_scalarnew_transposed_matrix(): + A = Matrix.from_coo([0, 1, 4], [1, 2, 6], [10.0, 20.0, 30.0], nrows=5, ncols=7) + AT = A.T + _assert_scalars_equal(AT[1, 0].new(), _new_slow(AT[1, 0])) + _assert_scalars_equal(AT[6, 4].new(), _new_slow(AT[6, 4])) + _assert_scalars_equal(AT[0, 0].new(), _new_slow(AT[0, 0])) + _assert_scalars_equal(AT[-1, -1].new(), _new_slow(AT[-1, -1])) + + +def test_scalarnew_name_honored(): + v = Vector.from_coo([0, 5, 9], [1.5, 2.5, 3.5], size=10) + s = v[5].new(name="myscalar") + assert s.name == "myscalar" + + +def test_scalarnew_dtype_cast_fallback(): + v = Vector.from_coo([0, 5, 9], [1.5, 2.5, 3.5], size=10) + _assert_scalars_equal(v[5].new(dtype=dtypes.INT32), _new_slow(v[5], dtype=dtypes.INT32)) + _assert_scalars_equal(v[3].new(dtype=dtypes.INT32), _new_slow(v[3], dtype=dtypes.INT32)) + + +def test_scalarnew_is_cscalar_fallback(): + v = Vector.from_coo([0, 5, 9], [1.5, 2.5, 3.5], size=10) + f_cs = v[5].new(is_cscalar=True) + s_cs = _new_slow(v[5], is_cscalar=True) + assert f_cs.is_cscalar + assert s_cs.is_cscalar + assert f_cs.value == s_cs.value == 2.5 + + +def test_scalarnew_udt_fallback(): + udt = gb.dtypes.register_anonymous( + np.dtype([("nx", np.int64), ("ny", np.float64)]), "NewFastPathProbe" + ) + u = Vector(udt, 3) + u[1] = (7, 2.5) + fu = u[1].new() + su = _new_slow(u[1]) + assert fu.dtype == su.dtype + assert fu.value["nx"] == su.value["nx"] == 7 + assert u[0].new().is_empty + assert _new_slow(u[0]).is_empty + + +def test_scalarnew_recorder_fallback(): + v = Vector.from_coo([0], [1.5], size=4) + with Recorder() as rec: + assert v[0].new().value == 1.5 + assert "extractElement" in "".join(rec.data) + + +def test_scalarnew_autocompute_value_path(): + """With autocompute on, ``v[i].value`` / ``float(v[i])`` still match ``get``.""" + v = Vector.from_coo([0, 5, 9], [1.5, 2.5, 3.5], size=10) + with gb.config.set(autocompute=True): + for idx in [0, 5, 9, 3]: # 3 is empty + want = v.get(idx) + got_val = v[idx].value + got_get = v[idx].get() + if want is None: + assert got_val is None, (idx, got_val, got_get) + assert got_get is None, (idx, got_val, got_get) + else: + assert got_val == want, (idx, got_val, got_get, want) + assert got_get == want, (idx, got_val, got_get, want) + assert float(v[5]) == 2.5 + assert int(v[0]) == 1 # 1.5 -> int truncates via Scalar.__int__ + + +def test_scalarnew_autocompute_false_raises(): + """With autocompute off, the value path raises but ``.new()`` still works.""" + v = Vector.from_coo([0, 5, 9], [1.5, 2.5, 3.5], size=10) + with gb.config.set(autocompute=False): + with pytest.raises(TypeError): + v[5].value + with pytest.raises(TypeError): + float(v[5]) + assert v[5].new().value == 2.5 + + +# --------------------------------------------------------------------------- +# ScalarIndexExpr value reads (.value / float / int / ...): the single-extract +# fast path in automethods._get_value via ScalarIndexExpr._extract_fast. +# --------------------------------------------------------------------------- +# The reference is the pre-fast-path resolution: _get_value used to resolve the +# expression to a GrB_Scalar (is_cscalar=False) via `.new()`, then read that +# scalar's attribute. `_new_slow` reproduces that GrB_Scalar exactly, so +# `getattr(_new_slow(expr), attr)` is the old behavior for every read attr. +_VALUE_READ_ATTRS = [ + "value", + "is_empty", + "_is_empty", + "__float__", + "__int__", + "__complex__", + "__bool__", + "__index__", + "__array__", +] + + +def _read_attr(scalar, attr): + """Fingerprint of a value-read attr on a Scalar (or index expr). + + Captures the whole access (``__index__`` raises via its property on + non-integral dtypes) and reduces the result to ``(repr, typename)`` so NaN + from an empty float scalar and numpy arrays compare structurally, not by + value. + """ + + def go(): + getter = getattr(scalar, attr) + if attr in ("__float__", "__int__", "__complex__", "__bool__", "__index__", "__array__"): + return getter() + return getter + + cap = _capture(go) + if cap[0] == "val": + return ("val", repr(cap[1]), type(cap[1]).__name__) + return cap + + +@pytest.mark.parametrize("dt", BUILTIN_DTYPES, ids=DTYPE_IDS) +def test_scalarvalue_dtype_parity(dt): + """Every read attr matches the pre-fast-path GrB_Scalar result, hit and miss.""" + val = _hit_value(dt) + v = Vector(dt, 10) + v[2] = val + A = Matrix(dt, 5, 7) + A[1, 2] = val + with gb.config.set(autocompute=True): + for probe in [v[2], v[3], A[1, 2], A[0, 0]]: + slow = _new_slow(probe) # GrB_Scalar, old path + for attr in _VALUE_READ_ATTRS: + assert _read_attr(probe, attr) == _read_attr(slow, attr), (dt.name, attr) + + +def test_scalarvalue_transposed_matrix(): + A = Matrix.from_coo([0, 1, 4], [1, 2, 6], [10.0, 20.0, 30.0], nrows=5, ncols=7) + AT = A.T + with gb.config.set(autocompute=True): + for r, c in [(1, 0), (6, 4), (0, 0), (-1, -1)]: + for attr in _VALUE_READ_ATTRS: + got = _read_attr(AT[r, c], attr) + want = _read_attr(_new_slow(AT[r, c]), attr) + assert got == want, (r, c, attr, got, want) + + +def test_scalarvalue_negative_and_numpy_index(): + v = Vector.from_coo([0, 5, 9], [1.5, 2.5, 3.5], size=10) + with gb.config.set(autocompute=True): + for idx in [-1, -5, -2, np.int64(5)]: + assert repr(v[idx].value) == repr(_new_slow(v[idx]).value) + got = _capture(lambda: float(v[idx])) # noqa: B023 + want = _capture(lambda: float(_new_slow(v[idx]))) # noqa: B023 + assert got == want, (idx, got, want) + + +def test_scalarvalue_udt_fallback(): + """UDT reads defer to `.new()`; Scalar.value's numpy conversion still runs.""" + udt = gb.dtypes.register_anonymous( + np.dtype([("vx", np.int64), ("vy", np.float64)]), "ValueFastPathProbe" + ) + u = Vector(udt, 3) + u[1] = (7, 2.5) + with gb.config.set(autocompute=True): + got = u[1].value + want = _new_slow(u[1]).value + assert got["vx"] == want["vx"] == 7 + assert got["vy"] == want["vy"] == 2.5 + assert u[0].value is None + assert u[0].is_empty is True + + +def test_scalarvalue_recorder_fallback(): + """With a Recorder active the value path still records an extractElement.""" + v = Vector.from_coo([0], [1.5], size=4) + with gb.config.set(autocompute=True), Recorder() as rec: + assert v[0].value == 1.5 + assert "extractElement" in "".join(rec.data) + + +def test_scalarvalue_pending_value(): + """A value written in non-blocking mode is read back correctly (iso/pending).""" + v = Vector(dtypes.FP64, 100) + v[3] = 2.25 + with gb.config.set(autocompute=True): + assert v[3].value == 2.25 + assert float(v[3]) == 2.25 + + +def test_scalarvalue_index_integral_only(): + """__index__ is available on integral dtypes and absent on float, via the fast path.""" + vi = Vector.from_coo([2], [7], dtype=dtypes.INT64, size=5) + vf = Vector.from_coo([2], [7.5], dtype=dtypes.FP64, size=5) + with gb.config.set(autocompute=True): + assert vi[2].__index__() == 7 + with pytest.raises(AttributeError): + vf[2].__index__ + + +# --------------------------------------------------------------------------- +# parse_index plain-int lane +# --------------------------------------------------------------------------- +def _resolved_fingerprint(expr): + """A comparable fingerprint of a resolved scalar index expression.""" + ax = expr.resolved_indexes.indices + return tuple((a.size, a.index.value, a.dimsize, a._carg.__class__.__name__) for a in ax) + + +def _index_probe_vector(v, i): + return _capture(lambda: (v[i].new().value, _resolved_fingerprint(v[i]))) + + +def _index_probe_matrix(A, r, c): + return _capture(lambda: (A[r, c].new().value, _resolved_fingerprint(A[r, c]))) + + +@pytest.mark.parametrize("i", [0, 2, 4, 9, -1, -10, -6]) +def test_indexparse_vector_valid(i): + """The plain-int fast lane matches the numpy-int lane for valid indices.""" + v = Vector.from_coo([0, 2, 4, 9], [10.0, 20.0, 30.0, 90.0], size=10) + fast = _index_probe_vector(v, i) + slow = _index_probe_vector(v, np.int64(i)) + assert fast == slow, f"fast={fast} slow={slow}" + + +@pytest.mark.parametrize("i", [10, 11, 1000, -11, -100]) +def test_indexparse_vector_out_of_range(i): + """Both lanes raise the identical IndexError for out-of-range indices.""" + v = Vector.from_coo([0, 2, 4, 9], [10.0, 20.0, 30.0, 90.0], size=10) + fast = _capture(lambda: v[i]) + slow = _capture(lambda: v[np.int64(i)]) + assert fast == slow, f"fast={fast} slow={slow}" + + +@pytest.mark.parametrize("b", [True, False]) +def test_indexparse_vector_bool(b): + """A bool must not take the int lane; it keeps the pre-existing TypeError.""" + v = Vector.from_coo([0, 2, 4, 9], [10.0, 20.0, 30.0, 90.0], size=10) + r = _capture(lambda: v[b]) + assert r[0] == "exc", r + assert r[1] == "TypeError", r + + +@pytest.mark.parametrize("bad", [2.0, "x", None]) +def test_indexparse_vector_non_integer(bad): + v = Vector.from_coo([0, 2, 4, 9], [10.0, 20.0, 30.0, 90.0], size=10) + r = _capture(lambda: v[bad]) + assert r[0] == "exc", r + + +@pytest.mark.parametrize("r", [0, 1, 4, -1, -5]) +@pytest.mark.parametrize("c", [0, 2, 6, -1, -7]) +def test_indexparse_matrix_valid(r, c): + A = Matrix.from_coo([0, 1, 4], [0, 2, 6], [1.0, 2.0, 3.0], nrows=5, ncols=7) + fast = _index_probe_matrix(A, r, c) + slow = _capture( + lambda: ( + A[np.int64(r), np.int64(c)].new().value, + _resolved_fingerprint(A[np.int64(r), np.int64(c)]), + ) + ) + assert fast == slow, f"fast={fast} slow={slow}" + + +@pytest.mark.parametrize("r", [1, -1]) +@pytest.mark.parametrize("c", [2, -2]) +def test_indexparse_matrix_mixed_lane(r, c): + """Plain int on one axis, numpy int on the other, must still agree.""" + A = Matrix.from_coo([0, 1, 4], [0, 2, 6], [1.0, 2.0, 3.0], nrows=5, ncols=7) + a = _capture(lambda: (A[r, np.int64(c)].new().value, _resolved_fingerprint(A[r, np.int64(c)]))) + b = _capture(lambda: (A[np.int64(r), c].new().value, _resolved_fingerprint(A[np.int64(r), c]))) + assert a == b, f"a={a} b={b}" + + +@pytest.mark.parametrize(("r", "c"), [(5, 0), (100, 0), (-6, 0), (0, 7), (0, 1000), (0, -8)]) +def test_indexparse_matrix_out_of_range(r, c): + A = Matrix.from_coo([0, 1, 4], [0, 2, 6], [1.0, 2.0, 3.0], nrows=5, ncols=7) + fast = _capture(lambda: A[r, c]) + slow = _capture(lambda: A[np.int64(r), np.int64(c)]) + assert fast == slow, f"fast={fast} slow={slow}" + + +@pytest.mark.parametrize("probe", ["row", "col", "false"]) +def test_indexparse_matrix_bool(probe): + A = Matrix.from_coo([0, 1, 4], [0, 2, 6], [1.0, 2.0, 3.0], nrows=5, ncols=7) + probes = { + "row": lambda: A[True, 0], + "col": lambda: A[0, True], + "false": lambda: A[False, 1], + } + r = _capture(probes[probe]) + assert r[0] == "exc", r + assert r[1] == "TypeError", r + + +def test_indexparse_fancy_unaffected(): + """List / ndarray fancy indexing is untouched by the plain-int lane.""" + v = Vector.from_coo([0, 2, 4, 9], [10.0, 20.0, 30.0, 90.0], size=10) + fast = _capture(lambda: v[[1, 3, 5]].new().to_coo()) + slow = _capture(lambda: v[np.array([1, 3, 5])].new().to_coo()) + assert fast[0] == slow[0] == "val" + # to_coo returns numpy arrays; compare structurally + (fi, fx), (si, sx) = fast[1], slow[1] + assert np.array_equal(fi, si) + assert np.array_equal(fx, sx) diff --git a/graphblas/tests/test_formatting.py b/graphblas/tests/test_formatting.py index a6522dcef..264dc0502 100644 --- a/graphblas/tests/test_formatting.py +++ b/graphblas/tests/test_formatting.py @@ -146,41 +146,26 @@ def t(): def test_no_pandas_repr(A, C, v, w): - # This is a bit of a hack... + # The rich repr is hand-rendered, so it no longer depends on pandas: with + # pandas marked absent the output is byte-identical to the pandas-present + # output, data grid and all. (When pandas is genuinely absent both branches + # use the same pandas-free path, so this still exercises that code.) + objs = [A, A.T, C, C.S, ~C.V, v, v.S, ~w.V, w] + expected = [repr(x) for x in objs] has_pandas_prev = formatting.has_pandas formatting.has_pandas = False try: - repr_printer(A, "A", indent=8) - assert repr(A) == ( - '"A_1" nvals nrows ncols dtype format\n' - "gb.Matrix 3 1 5 INT64 bitmapr" - ) - repr_printer(A.T, "A.T", indent=8) - assert repr(A.T) == ( - '"A_1.T" nvals nrows ncols dtype format\n' - "gb.TransposedMatrix 3 5 1 INT64 bitmapc" - ) - repr_printer(C.S, "C.S", indent=8) - assert repr(C.S) == ( - '"C.S" nvals nrows ncols dtype format\n' - "StructuralMask\n" - "of gb.Matrix 8 70 77 INT64 hypercsr" - ) - repr_printer(v, "v", indent=8) - assert repr(v) == ( - '"v" nvals size dtype format\ngb.Vector 3 5 FP64 bitmap' - ) - repr_printer(~w.V, "~w.V", indent=8) - assert repr(~w.V) == ( - '"~w.V" nvals size dtype format\n' - "ComplementedValueMask\n" - "of gb.Vector 4 77 INT64 bitmap" - ) + actual = [repr(x) for x in objs] finally: formatting.has_pandas = has_pandas_prev + assert actual == expected + # The data grid is rendered, not just the header: the border line and the + # values are present. + lines = repr(A).split("\n") + assert lines[2].startswith("----") + assert lines[-1] == "0 0 1 2" -@pytest.mark.skipif("not pd") def test_matrix_repr_small(A, B): repr_printer(A, "A") assert repr(A) == ( @@ -212,7 +197,6 @@ def test_matrix_repr_small(A, B): ) -@pytest.mark.skipif("not pd") def test_matrix_mask_repr_small(A): repr_printer(A.S, "A.S") assert repr(A.S) == ( @@ -404,7 +388,6 @@ def test_matrix_mask_repr_large(C): ) -@pytest.mark.skipif("not pd") def test_vector_repr_small(v): repr_printer(v, "v") assert repr(v) == ( @@ -429,7 +412,6 @@ def test_vector_repr_large(w): ) -@pytest.mark.skipif("not pd") def test_vector_mask_repr_small(v): repr_printer(v.S, "v.S") assert repr(v.S) == ( @@ -517,140 +499,23 @@ def test_scalar_repr(s, t): def test_no_pandas_repr_html(A, C, v, w): - # This is a bit of a hack... + # _repr_html_ is hand-rendered too: marking pandas absent yields output that + # is byte-identical to the pandas-present output, data table included. + objs = [A, A.T, C, C.S, ~C.V, v, v.S, ~w.V, w] + expected = [repr_html(x) for x in objs] has_pandas_prev = formatting.has_pandas formatting.has_pandas = False try: - html_printer(A, "A", indent=8) - assert repr_html(A) == ( - "
" - f"{CSS_STYLE}" - '
A1
\n' - '\n' - " \n" - ' \n' - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - "
gb.Matrix
nvals
nrows
ncols
dtype
format
315INT64bitmapr
\n" - "
\n" - "(Install pandas to see a preview of the data)
" - ) - html_printer(A.T, "A.T", indent=8) - assert repr_html(A.T) == ( - "
" - f"{CSS_STYLE}" - '
A1.T
\n' - '\n' - " \n" - ' \n' - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - "
gb.TransposedMatrix
nvals
nrows
ncols
dtype
format
351INT64bitmapc
\n" - "
\n" - "(Install pandas to see a preview of the data)
" - ) - html_printer(C.S, "C.S", indent=8) - assert repr_html(C.S) == ( - "
" - f"{CSS_STYLE}" - '
C.S
\n' - '\n' - " \n" - ' \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - "
StructuralMask\n'
-            "of\n"
-            "gb.Matrix
nvals
nrows
ncols
dtype
format
87077INT64hypercsr
\n" - "
\n" - "(Install pandas to see a preview of the data)
" - ) - html_printer(v, "v", indent=8) - assert repr_html(v) == ( - "
" - f"{CSS_STYLE}" - '
v
\n' - '\n' - " \n" - ' \n' - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - "
gb.Vector
nvals
size
dtype
format
35FP64bitmap
\n" - "
\n" - "(Install pandas to see a preview of the data)
" - ) - html_printer(~w.V, "~w.V", indent=8) - assert repr_html(~w.V) == ( - "
" - f"{CSS_STYLE}" - '
~w.V
\n' - '\n' - " \n" - ' \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - " \n" - "
ComplementedValueMask\n'
-            "of\n"
-            "gb.Vector
nvals
size
dtype
format
477INT64bitmap
\n" - "
\n" - "(Install pandas to see a preview of the data)
" - ) + actual = [repr_html(x) for x in objs] finally: formatting.has_pandas = has_pandas_prev + assert actual == expected + # The data table is rendered, not the "install pandas" placeholder. + html = repr_html(A) + assert "install" not in html.lower() + assert '' in html -@pytest.mark.skipif("not pd") def test_matrix_repr_html_small(A, B): html_printer(A, "A") assert repr_html(A) == ( @@ -845,7 +710,6 @@ def test_matrix_repr_html_small(A, B): ) -@pytest.mark.skipif("not pd") def test_matrix_mask_repr_html_small(A): html_printer(A.S, "A.S") assert repr_html(A.S) == ( @@ -2121,7 +1985,6 @@ def test_matrix_mask_repr_html_large(C): ) -@pytest.mark.skipif("not pd") def test_vector_repr_html_small(v): html_printer(v, "v") assert repr_html(v) == ( @@ -2267,7 +2130,6 @@ def test_vector_repr_html_large(w): ) -@pytest.mark.skipif("not pd") def test_vector_mask_repr_html_small(v): html_printer(v.S, "v.S") assert repr_html(v.S) == ( @@ -2889,7 +2751,6 @@ def test_apply_repr(v): ) -@pytest.mark.skipif("not pd") def test_apply_repr_html(v): html_printer(v.apply(unary.one), "v.apply(unary.one)") assert repr_html(v.apply(unary.one)) == ( @@ -2922,7 +2783,6 @@ def test_mxm_repr(A, B): ) -@pytest.mark.skipif("not pd") def test_mxm_repr_html(A, B): html_printer(A.mxm(B), "A.mxm(B)") assert repr_html(A.mxm(B)) == ( @@ -2957,7 +2817,6 @@ def test_mxv_repr(A, v): ) -@pytest.mark.skipif("not pd") def test_mxv_repr_html(A, v): html_printer(A.mxv(v), "A.mxv(v)") assert repr_html(A.mxv(v)) == ( @@ -2980,7 +2839,6 @@ def test_mxv_repr_html(A, v): ) -@pytest.mark.skipif("not pd") def test_matrix_reduce_columns_repr_html(A): # This is implemented using the transpose of A, so make sure we're oriented correctly! html_printer(A.reduce_columnwise(), "A.reduce_columnwise()") @@ -3014,7 +2872,6 @@ def test_matrix_reduce_repr(C, v): ) -@pytest.mark.skipif("not pd") def test_matrix_reduce_repr_html(C, v): html_printer(C.reduce_scalar(), "C.reduce_scalar()", indent=8) assert repr_html(C.reduce_scalar()) == ( @@ -3035,7 +2892,6 @@ def test_matrix_reduce_repr_html(C, v): ) -@pytest.mark.skipif("not pd") def test_matrix_huge(): M = Matrix(int, nrows=2**60, ncols=2**60, name="M") repr_printer(M, "M") @@ -3061,7 +2917,6 @@ def test_matrix_huge(): assert M.isequal(M2) -@pytest.mark.skipif("not pd") def test_matrix_huge_html(): M = Matrix(int, nrows=2**60, ncols=2**60, name="M") html_printer(M, "M") @@ -3254,7 +3109,6 @@ def test_matrix_huge_html(): ) -@pytest.mark.skipif("not pd") def test_vector_huge(): v = Vector(int, size=2**60) repr_printer(v, "v") @@ -3269,7 +3123,6 @@ def test_vector_huge(): assert v2.isequal(v) -@pytest.mark.skipif("not pd") def test_vector_huge_html(): v = Vector(int, size=2**60) html_printer(v, "v") @@ -3407,7 +3260,6 @@ def test_vector_huge_html(): ) -@pytest.mark.skipif("not pd") def test_sparse_vector_repr(): v = Vector.from_coo([100 * i for i in range(100)], [10 * i for i in range(100)], name="v") repr_printer(v, "v") @@ -3561,7 +3413,6 @@ def test_sparse_vector_repr(): ) -@pytest.mark.skipif("not pd") def test_sparse_matrix_repr(): A = Matrix.from_coo( [100 * i for i in range(100)], [10 * i for i in range(100)], list(range(100)), name="A" @@ -3735,7 +3586,6 @@ def test_sparse_matrix_repr(): ) -@pytest.mark.skipif("not pd") def test_infix_expr_repr_html(A, B, v): html_printer(v & v, "v & v") assert repr_html(v & v) == ( @@ -3934,7 +3784,6 @@ def test_infix_expr_repr_html(A, B, v): ) -@pytest.mark.skipif("not pd") def test_infix_expr_repr(A, B, v): repr_printer(v & v, "v & v") assert repr(v & v) == ( @@ -4010,7 +3859,6 @@ def test_infix_expr_repr(A, B, v): ) -@pytest.mark.skipif("not pd") def test_inner_outer_repr_html(v): html_printer(v.inner(v), "v.inner(v)") assert repr_html(v.inner(v)) == ( @@ -4052,7 +3900,6 @@ def test_inner_outer_repr_html(v): ) -@pytest.mark.skipif("not pd") def test_inner_outer_repr(v): # XXX: hmm, having `(GrB_Matrix)` here isn't so pretty repr_printer(v.inner(v), "v.inner(v)") @@ -4073,8 +3920,6 @@ def test_inner_outer_repr(v): @autocompute def test_autocompute(A, B, v): - if not pd: # pragma: no cover (import) - pytest.skip("needs pandas") repr_printer(A & A, "A & A") assert repr(A & A) == ( "gb.MatrixEwiseMultExpr nrows ncols left_dtype right_dtype\n" @@ -4155,8 +4000,6 @@ def test_autocompute(A, B, v): @autocompute def test_autocompute_html(A, B, v): - if not pd: # pragma: no cover (import) - pytest.skip("needs pandas") html_printer(A & A, "A & A") assert repr_html(A & A) == ( "
" @@ -4478,7 +4321,6 @@ def test_autocompute_html(A, B, v): ) -@pytest.mark.skipif("not pd") def test_display_nan(): v = Vector.from_coo([0, 1], [1.0, np.nan], size=3, name="v") repr_printer(v, "v") @@ -4612,7 +4454,6 @@ def test_display_nan(): ) -@pytest.mark.skipif("not pd") def test_large_iso(): A = Matrix(int, nrows=2**60, ncols=2**60) A[:, :] << 1 @@ -4841,7 +4682,6 @@ def test_index_expr_matrix_html(A): ) -@pytest.mark.skipif("not pd") def test_scalar_as_vector(): s = Scalar.from_value(5, is_cscalar=False) # pragma: is_grbscalar v = s._as_vector() @@ -4900,8 +4740,6 @@ def test_scalar_as_vector(): @autocompute def test_index_expr_autocompute(v): - if not pd: # pragma: no cover (import) - pytest.skip("needs pandas") html_printer(v[[0, 1]], "v[[0, 1]]") assert repr_html(v[[0, 1]]) == ( "
" @@ -4955,7 +4793,6 @@ def test_index_expr_autocompute(v): ) -@pytest.mark.skipif("not pd") def test_udt(): record_dtype = np.dtype([("x", np.bool_), ("y", np.int64)], align=True) udt = dtypes.register_anonymous(record_dtype, "record_dtype") @@ -5005,7 +4842,6 @@ def test_udt(): ) -@pytest.mark.skipif("not pd") def test_empty(): v = Vector(int, 0) repr_printer(v, "v") @@ -5030,7 +4866,6 @@ def test_empty(): ) -@pytest.mark.skipif("not pd") def test_vector_as_matrix(): v = Vector.from_coo([1], [2], name="v_A") A = v._as_matrix() diff --git a/graphblas/tests/test_indexbinary.py b/graphblas/tests/test_indexbinary.py index c56bb86f9..82ffbdbf0 100644 --- a/graphblas/tests/test_indexbinary.py +++ b/graphblas/tests/test_indexbinary.py @@ -300,6 +300,32 @@ def test_bind_raw_array_udt_theta(): delattr(indexbinary, "raw_array_udt_op") +def test_array_udt_udf_shape_is_checked(): + """An IBO whose UDF builds a wrong-shape array is rejected at registration. + + The same guard the binary and unary paths get. Numba's ``Array`` type + carries ``ndim`` but not extents, so the mismatch is not a type error; left + to the wrapper it raises inside the cfunc, where Numba prints the traceback + and returns, leaving the element as SuiteSparse found it. + """ + # A shape of its own: ``register_anonymous`` caches by numpy dtype, so + # sharing one with another test would hand back that test's DataType. + arr_udt = dtypes.register_anonymous(np.dtype((np.float64, (12,))), "_IboShapeArr12") + + def _truncate(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return (x + y)[:2] + + op = indexbinary.register_anonymous(_truncate, is_udt=True) + with pytest.raises(UdfParseError, match=r"shape \(2,\) when run on sample values"): + op[arr_udt] + + def _combine(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x + y + theta + + good = indexbinary.register_anonymous(_combine, is_udt=True) + assert good[arr_udt].return_type is arr_udt + + def test_bind_raw_udt_theta_without_dtype_errors(): """A raw UDT theta with no dtype can't be inferred; the error is clear and actionable.""" indexbinary.register_new("raw_no_dtype_op", _ibo_return_x, is_udt=True) diff --git a/graphblas/tests/test_io.py b/graphblas/tests/test_io.py index a8d80e9ca..e5eb0020d 100644 --- a/graphblas/tests/test_io.py +++ b/graphblas/tests/test_io.py @@ -160,6 +160,113 @@ def test_matrix_to_from_networkx(): assert M.shape == (1, 1) +@pytest.mark.skipif("not nx") +def test_from_networkx_undirected(): + # Undirected graphs go through from_networkx's direct COO path (no scipy), + # which symmetrizes off-diagonal entries and keeps self-loops single-counted. + G = nx.Graph() + G.add_weighted_edges_from([(0, 1, 2.0), (0, 0, 7.0), (1, 2, 3.0)]) + M = gb.io.from_networkx(G) + expected = gb.Matrix.from_coo([0, 0, 1, 1, 2], [0, 1, 0, 2, 1], [7.0, 2.0, 2.0, 3.0, 3.0]) + assert M.isequal(expected, check_dtype=True) + + # weight=None ignores edge weights (all entries 1) and yields an int Matrix + M_none = gb.io.from_networkx(G, weight=None) + expected_none = gb.Matrix.from_coo([0, 0, 1, 1, 2], [0, 1, 0, 2, 1], 1) + assert M_none.isequal(expected_none, check_dtype=True) + + +@pytest.mark.skipif("not nx or not ss") +@pytest.mark.parametrize( + "graph_cls", [nx.Graph, nx.DiGraph, nx.MultiGraph, nx.MultiDiGraph] if nx else [] +) +def test_from_networkx_rejects_array_weights(graph_cls): + # Sequence weights of uniform length infer a 2-D numeric array, which passes + # a dtype-kind test but which from_coo would happily read as a UDT. scipy has + # always rejected these, so the direct path must defer rather than quietly + # widen what from_networkx accepts. + G = graph_cls() + G.add_edge(0, 1, weight=[1, 2]) + G.add_edge(1, 0, weight=[3, 4]) + with pytest.raises(ValueError, match="must be 1-D"): + gb.io.from_networkx(G) + + +@pytest.mark.skipif("not nx or not ss") +@pytest.mark.parametrize( + "graph_cls", [nx.Graph, nx.DiGraph, nx.MultiGraph, nx.MultiDiGraph] if nx else [] +) +def test_from_networkx_rejects_string_weights(graph_cls): + # String weights infer a 1-D weights summed + G.add_edge(0, 0, weight=2.5) # parallel self-loop -> diagonal summed + G.add_edge(3, 4) # missing weight attr -> default 1 + G.add_node(9) # isolated node + + A = nx.to_scipy_sparse_array(G, weight="weight") + reference = gb.io.from_scipy_sparse(A) + M = gb.io.from_networkx(G, weight="weight") + assert M.isequal(reference, check_dtype=True) + assert M.shape == reference.shape + + +@pytest.mark.skipif("not nx") +def test_from_networkx_multigraph_is_scipy_free(monkeypatch): + # A numeric-weight multigraph must ingest via the direct coo path, not the + # scipy fallback, so networkx ingest no longer requires scipy for this case. + import graphblas.io._networkx as _gnx + + def _boom(*args, **kwargs): # pragma: no cover (only runs if the direct path regresses) + raise AssertionError("scipy fallback should not be used for a numeric multigraph") + + monkeypatch.setattr(_gnx, "_from_networkx_via_scipy", _boom) + G = nx.MultiDiGraph() + G.add_weighted_edges_from([(0, 1, 2.0), (0, 1, 3.0), (2, 0, 1.0), (1, 1, 4.0), (1, 1, 0.5)]) + M = gb.io.from_networkx(G) + assert M[0, 1].new().value == 5.0 # parallel edges summed + assert M[1, 1].new().value == 4.5 # parallel self-loops summed + + @pytest.mark.skipif("not ss") @pytest.mark.parametrize("engine", ["auto", "scipy", "fmm"]) def test_mmread_mmwrite(engine): @@ -388,6 +495,24 @@ def test_matrix_market_bad_engine(): gb.io.mmread(mm_out, engine="bad_engine") +@pytest.mark.skipif("not ss") +@pytest.mark.skipif("fmm is not None") +def test_matrix_market_fmm_engine_unavailable(): + # Naming the deprecated engine warns before anything else, so the caller hears + # about the deprecation even when the install that would satisfy it is missing. + A = gb.Matrix.from_coo([0, 0, 3, 5], [1, 4, 0, 2], [1, 0, 2, -1], nrows=7, ncols=6) + with ( + pytest.warns(DeprecationWarning, match="fast_matrix_market is no longer maintained"), + pytest.raises(ImportError, match="required to write Matrix Market files"), + ): + gb.io.mmwrite(BytesIO(), A, engine="fmm") + with ( + pytest.warns(DeprecationWarning, match="fast_matrix_market is no longer maintained"), + pytest.raises(ImportError, match="required to read Matrix Market files"), + ): + gb.io.mmread(BytesIO(), engine="fast_matrix_market") + + @pytest.mark.skipif("not ss") def test_scipy_sparse(): a = np.arange(12).reshape(3, 4) diff --git a/graphblas/tests/test_matrix.py b/graphblas/tests/test_matrix.py index 4befea2cf..3176a310e 100644 --- a/graphblas/tests/test_matrix.py +++ b/graphblas/tests/test_matrix.py @@ -1235,6 +1235,31 @@ def test_apply_indexunary(A): assert pickle.loads(pickle.dumps(indexunary.tril[int])) is indexunary.tril[int] +def test_indexunary_helpers(A): + # indexunary.value/row/column mirror select's helpers (GH #239), but they + # produce apply expressions (BOOL) rather than select expressions. + assert indexunary.value(A > 3).new().isequal(A.apply(indexunary.valuegt, 3).new()) + assert indexunary.value(A == 3).new().isequal(A.apply(indexunary.valueeq, 3).new()) + # Only rowle/rowgt and colle/colgt exist, so `<` and `>=` are rewritten with + # a thunk shift, and `==`/`!=` have no counterpart (mirrors select's helpers). + assert indexunary.row(A <= 2).new().isequal(A.apply(indexunary.rowle, 2).new()) + assert indexunary.row(A < 3).new().isequal(A.apply(indexunary.rowle, 2).new()) + assert indexunary.row(A >= 3).new().isequal(A.apply(indexunary.rowgt, 2).new()) + assert indexunary.row(A > 2).new().isequal(A.apply(indexunary.rowgt, 2).new()) + assert indexunary.column(A < 3).new().isequal(A.apply(indexunary.colle, 2).new()) + assert indexunary.column(A > 2).new().isequal(A.apply(indexunary.colgt, 2).new()) + for expr in [indexunary.value(A > 3), indexunary.row(A <= 2), indexunary.column(A < 3)]: + assert expr.new().dtype == dtypes.BOOL + with pytest.raises(TypeError, match="indexunary.value"): + indexunary.value(A) + with pytest.raises(TypeError, match="indexunary.row"): + indexunary.row(A | A) + with pytest.raises(ValueError, match="roweq"): + indexunary.row(A == 3) + with pytest.raises(ValueError, match="coleq"): + indexunary.column(A == 3) + + def test_select(A): A3 = Matrix.from_coo([0, 3, 3, 6], [3, 0, 2, 4], [3, 3, 3, 3], nrows=7, ncols=7) w1 = A.select(select.valueeq, 3).new() @@ -2929,6 +2954,7 @@ def test_expr_is_like_matrix(A): "__call__", "__del__", "__delitem__", + "__getattr__", "__lshift__", "__setitem__", "_assign_element", @@ -2995,6 +3021,7 @@ def test_index_expr_is_like_matrix(A): expected = { "__del__", "__delitem__", + "__getattr__", "__setitem__", "_assign_element", "_delete_element", @@ -4545,3 +4572,152 @@ def test_setdiag(): A.setdiag(30, mask=v.S) expected[0, 0] = 30 assert A.isequal(expected) + + +def test_constructor_rejects_arraylike_first_arg(): + # The first positional arg is the dtype; passing data instead used to raise a + # confusing "Unknown dtype" ValueError. It should now raise TypeError pointing + # at the from_* constructors. + with pytest.raises(TypeError, match="Matrix.*dtype.*from_coo.*from_dense"): + Matrix([[1, 2], [3, 4]]) + with pytest.raises(TypeError, match="Matrix.*expects a dtype"): + Matrix(np.zeros((2, 2))) + # Valid dtype-first signatures still work, including list/tuple dtype specs + assert Matrix(int, 2, 2).dtype == dtypes.INT64 + assert Matrix("INT64", nrows=2, ncols=2).dtype == dtypes.INT64 + assert Matrix([("x", "i8"), ("y", "f8")], nrows=2, ncols=2).dtype._is_udt + assert Matrix((np.int32, (2, 2)), nrows=2, ncols=2).dtype._is_udt + # A non-array-like bad dtype keeps the original ValueError + with pytest.raises(ValueError, match="Unknown dtype"): + Matrix("not_a_dtype", nrows=2, ncols=2) + + +def test_wrong_kind_mask_on_matrix_raises(): + # A Vector mask on a Matrix output used to leak a raw cffi/C-signature + # error ("initializer for ctype 'struct GB_Matrix_opaque'"). It should + # now raise a clear TypeError, mirroring the Vector-output guard. + A = Matrix(int, 3, 3) + A[0, 0] = 1 + v = Vector(int, 3) + v[0] = 1 + # update path (full-matrix op with a Vector mask) + C = Matrix(int, 3, 3) + with pytest.raises(TypeError, match="Mask object must be type Matrix"): + C(mask=v.S) << A.ewise_mult(A) + # extract path + with pytest.raises(TypeError, match="Mask object must be type Matrix"): + A[:, :].new(mask=v.S) + # Valid Matrix mask on Matrix output still works + C = Matrix(int, 3, 3) + C(mask=A.S) << A.ewise_mult(A) + assert C.nvals == A.nvals + # Valid Vector input_mask broadcast on a Matrix extract still works + m = Vector(bool, 3) + m[0] = True + m[2] = True + assert A[0, [0, 1, 2]].new(input_mask=m.S) is not None + + +def test_wrong_kind_mask_on_dup_and_full_assign(): + # dup and whole-object update assign with `C(mask)[...] = A`, so a + # Vector mask used to reach GrB_Matrix_assign and leak the same raw + # cffi error as full-matrix operations. + A = Matrix(int, 3, 3) + A[0, 0] = 1 + v = Vector(bool, 3) + v[0] = True + v[1] = True + err = "Unable to use Vector mask on Matrix assignment to a Matrix" + with pytest.raises(TypeError, match=err): + A.dup(mask=v.S) + C = Matrix(int, 3, 3) + with pytest.raises(TypeError, match=err): + C(v.S) << A + with pytest.raises(TypeError, match=err): + C(v.V) << A + with pytest.raises(TypeError, match=err): + C(~v.S) << A + with pytest.raises(TypeError, match=err): + C(v.S)[:, :] << A + with pytest.raises(TypeError, match=err): + C[:, :](v.S) << A + # A Matrix mask on these paths still works + assert A.dup(mask=A.S).isequal(A) + C(A.S) << A + assert C.isequal(A) + # A Vector mask on a Matrix row/column assignment is still valid + w = Vector(int, 3) + w[0] = 10 + w[2] = 30 + C = Matrix(int, 3, 3) + C(v.S)[0, :] << w + assert C.isequal(Matrix.from_coo([0], [0], [10], nrows=3, ncols=3)) + C.clear() + C(v.S)[:, 0] << w + assert C.isequal(Matrix.from_coo([0], [0], [10], nrows=3, ncols=3)) + C.clear() + C[0, :](v.S) << w + assert C.isequal(Matrix.from_coo([0], [0], [10], nrows=3, ncols=3)) + C.clear() + C[:, 0](v.S) << w + assert C.isequal(Matrix.from_coo([0], [0], [10], nrows=3, ncols=3)) + + # The guard keys on dimension, not exact type, so a Vector subclass's + # mask is caught too instead of leaking the raw cffi error + class _VecSub(Vector): + pass + + mv = _VecSub(bool, 3) + mv[0] = True + with pytest.raises(TypeError, match="Unable to use Vector mask"): + C(mv.S) << A + + +def test_new_constructor_misuse_hint(): + # `.new()` resolves expressions; it is not a constructor or an instance + # method on concrete objects. Both misuses should hint the right API. + A = Matrix(int, 3, 3) + v = Vector(int, 3) + s = Scalar.from_value(5) + # Class-level access stays a plain AttributeError: hinting it would need + # a metaclass, and both metaclass choices break something (see + # test_abc_mixin_subclass). + for cls in (Matrix, Vector, Scalar): + with pytest.raises(AttributeError, match="has no attribute 'new'"): + cls.new + # Instance-level misuse: A.new() -> hint .dup() + for obj in (A, v, s): + with pytest.raises(AttributeError, match=r"has no attribute 'new'.*\.dup\(\)"): + obj.new + # A genuinely-missing attribute keeps the plain AttributeError (no hint) + with pytest.raises(AttributeError, match="has no attribute 'frobnicate'"): + A.frobnicate + assert not hasattr(A, "new") + assert not hasattr(Matrix, "new") + + +def test_abc_mixin_subclass(): + # BaseType must stay metaclass-free: a plain `type` hint metaclass broke + # abc-based mixins with "metaclass conflict", and an ABCMeta-derived one + # leaks metaclass attributes into dir(Matrix), which the + # expression-surface guards report as drift. + import collections.abc + + class SizedMatrix(Matrix, collections.abc.Sized): + def __len__(self): + return 1 + + assert issubclass(SizedMatrix, collections.abc.Sized) + + +def test_transpose_hint(): + # A.transpose() should point at the .T property rather than fail with a + # bare "no attribute" message; transpose is not a method here. + A = Matrix(int, 3, 3) + v = Vector(int, 3) + with pytest.raises(AttributeError, match=r"has no attribute 'transpose'.*\.T"): + A.transpose + with pytest.raises(AttributeError, match=r"has no attribute 'transpose'.*\.T"): + v.transpose + # .T still works + assert A.T.shape == (3, 3) diff --git a/graphblas/tests/test_op.py b/graphblas/tests/test_op.py index ef5227f07..68699fe11 100644 --- a/graphblas/tests/test_op.py +++ b/graphblas/tests/test_op.py @@ -1,4 +1,8 @@ import itertools +import os +import subprocess +import sys +from pathlib import Path import numpy as np import pytest @@ -1155,6 +1159,15 @@ def test_from_string(): with pytest.raises(ValueError, match="Unknown agg string"): agg.from_string("bad_agg") + assert select.from_string("tril") is select.tril + assert select.from_string(">=") is select.valuege + assert indexunary.from_string("rowindex") is indexunary.rowindex + assert indexunary.from_string("rowindex[int]") is indexunary.rowindex[int] + + # Every namespace's from_string carries a docstring (GH #513) + for ns in [unary, binary, monoid, semiring, select, indexunary, agg, op]: + assert ns.from_string.__doc__ + @pytest.mark.skipif("not supports_udfs") @pytest.mark.slow @@ -1359,6 +1372,500 @@ def _this_or_that(val, idx, _, thunk): # pragma: no cover (numba) assert result.isequal(w) +@pytest.mark.skipif("not supports_udfs") +def test_udf_division_by_zero_follows_numpy(): + """Dividing by zero in a UDF returns numpy's answer instead of losing the element. + + Under Numba's default error model the division raises ZeroDivisionError + inside the cfunc, where Numba prints the traceback and returns, so + GraphBLAS keeps whatever was in the output element (in practice the + previous element's value). ``error_model="numpy"`` fixes that, but only if + it is set on the ``njit`` Dispatcher: a Dispatcher holds one compilation + per signature, and ``_build`` calls ``.compile(sig)`` before the wrapper + exists, so setting it on the ``cfunc`` alone comes too late to matter. + """ + + def _idiv(x, y): # pragma: no cover (numba) + return x // y + + op = BinaryOp.register_anonymous(_idiv, "_udf_zero_idiv") + v = Vector.from_coo([0, 1], [10, 20], dtype=dtypes.INT64) + w = Vector.from_coo([0, 1], [2, 0], dtype=dtypes.INT64) + assert op(v & w).new().to_coo()[1].tolist() == [5, 0] + + def _tdiv(x, y): # pragma: no cover (numba) + return x / y + + op = BinaryOp.register_anonymous(_tdiv, "_udf_zero_tdiv") + v = Vector.from_coo([0, 1], [1.0, 1.0], dtype=dtypes.FP64) + w = Vector.from_coo([0, 1], [2.0, 0.0], dtype=dtypes.FP64) + assert op(v & w).new().to_coo()[1].tolist() == [0.5, float("inf")] + + # Same guarantee for a UDT UDF, which reaches the cfunc by another route. + udt = dtypes.register_anonymous( + np.dtype([("dz_a", np.int64), ("dz_b", np.int64)], align=True), "_UdfDivZeroRec" + ) + + def _rec_idiv(x, y): # pragma: no cover (numba) + return (x["dz_a"] // y["dz_a"], x["dz_b"]) + + op = BinaryOp.register_anonymous(_rec_idiv, "_udf_zero_rec", is_udt=True) + v = Vector(udt, size=1) + v[0] = (10, 5) + w = Vector(udt, size=1) + w[0] = (0, 1) + got = v.ewise_mult(w, op).new()[0].new().value + assert (got["dz_a"], got["dz_b"]) == (0, 5) + + +@pytest.mark.skipif("not supports_udfs") +def test_select_op_outlives_source_indexunary(): + """A SelectOp keeps alive the IndexUnaryOp whose GraphBLAS handle it borrows. + + ``SelectOp._from_indexunary`` reuses the IndexUnaryOp's ``gb_obj`` rather + than allocating a second one, and ``register_anonymous`` drops that + IndexUnaryOp on the way out. Without an explicit reference the handle is + freed as soon as it is collected, and every use of the SelectOp raises + ``UninitializedObject``. + """ + import gc + + def _ne_thunk(x, i, j, thunk): # pragma: no cover (numba) + return x != thunk + + sel = SelectOp.register_anonymous(_ne_thunk) + gc.collect() + v = Vector.from_coo([0, 1, 2], [1, 5, 9]) + assert v.select(sel, 5).new().isequal(Vector.from_coo([0, 2], [1, 9], size=3)) + + +_jit_can_compile_cache = [] + + +def _jit_can_compile(): + """True when SuiteSparse has a C compiler it can actually use. + + Without one it falls back to the Numba cfunc and says nothing, so the + ``jit`` parameter would run the cfunc and report a pass. Any repair of + conda-baked compiler paths has already happened in + ``_auto_fix_jit_at_import``; calling ``fix_jit_config`` again from here + would rewrite process-wide compiler settings that no fixture restores. + + ``jit_compiler_is_usable`` alone is not enough: it only checks that the + configured compiler path exists on disk, and a runner can have the file + yet fail every compile (broken toolchain, missing headers). The + import-time probe already did a real compile, and SuiteSparse demotes + ``jit_c_control`` from ``'on'`` when that compile fails, so a control + still ``'on'`` here is the probe's success flag; ``test_ssjit`` keys its + skips on the same signal. The fixture calls this before it mutates the + control, and the cache keeps later per-test mutations from flipping it. + """ + if not _jit_can_compile_cache: + _jit_can_compile_cache.append( + gb.ss.jit_compiler_is_usable() and gb.ss.config["jit_c_control"] == "on" + ) + return _jit_can_compile_cache[0] + + +@pytest.fixture(params=["jit", "cfunc"]) +def udt_op_path(request): + """Pin SuiteSparse to one execution path for built-in UDT operators. + + Each auto-lifted UDT op carries both a JIT C definition and a Numba + cfunc, and SuiteSparse chooses between them per call depending on + whether a C compiler is available. A machine with one and a machine + without therefore run different code, so results have to hold on both. + """ + path = request.param + if backend != "suitesparse" or "jit_c_control" not in gb.ss.config: + if path == "jit": + pytest.skip("no SuiteSparse JIT on this backend") + yield path + return + previous = gb.ss.config["jit_c_control"] + if path == "jit": + if not _jit_can_compile(): + pytest.skip("JIT compilation not available (probe failed or compiler missing)") + # Set it rather than assume it. SuiteSparse demotes ``on`` to ``load`` + # after a failed compile, and a demoted control routes to the cfunc + # silently, so this parameter would pass while running the other path. + gb.ss.config["jit_c_control"] = "on" + else: + gb.ss.config["jit_c_control"] = "off" + try: + yield path + finally: + # Read before restoring: a demotion during the test is the signal that + # the kernel never compiled, which no assertion in the test can see. + demoted = path == "jit" and gb.ss.config["jit_c_control"] != "on" + gb.ss.config["jit_c_control"] = previous + if demoted: + pytest.fail("SuiteSparse demoted jit_c_control; the JIT path did not run") + + +def _udt_vectors(udt, xs, ys=None): + """Build one or two dense UDT vectors whose leaves all hold the given values. + + Values that repeat down the whole vector make it iso-valued, and + SuiteSparse answers those from a single element without reaching for a + JIT kernel, so callers pass varied data. + """ + names = udt.np_type.names + out = [] + for vals in (xs, ys): + if vals is None: + continue + v = Vector(udt, size=len(vals)) + for i, val in enumerate(vals): + v[i] = tuple(val for _ in names) if names else np.full(udt.np_type.subdtype[1], val) + out.append(v) + return out + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_mixed_record_dtypes_use_each_operands_own_dtype(udt_op_path): + """Two records sharing field names but not field types promote to the wider one. + + ``_check_udt_pair`` matches record operands on field names only, so their + leaf dtypes can differ. Offering the return-type resolver just the left + operand left it no choice but that record, so an int record over a float + record came back as the int record: ``7 / 2.0`` landed as 3 and ``6 / 0.0`` + as INT64_MAX. Swapping the operands changed the answer for the same pair. + """ + int_udt = dtypes.register_anonymous(np.dtype([("mxd_a", np.int64)], align=True), "_MixedRecInt") + float_udt = dtypes.register_anonymous( + np.dtype([("mxd_a", np.float64)], align=True), "_MixedRecFloat" + ) + v = Vector(int_udt, size=2) + v[0] = (6,) + v[1] = (7,) + w = Vector(float_udt, size=2) + w[0] = (0.0,) + w[1] = (2.0,) + + result = v.ewise_mult(w, binary.truediv).new() + assert result.dtype == float_udt, "result should promote to the float record" + assert result[0].new().value["mxd_a"] == float("inf") + assert result[1].new().value["mxd_a"] == 3.5 + + # The same pair the other way round must agree, which it did not when the + # resolver only ever saw the left operand. + swapped = w.ewise_mult(v, binary.truediv).new() + assert swapped.dtype == float_udt + assert swapped[1].new().value["mxd_a"] == 2.0 / 7.0 + + +def _bitwise_eq(got, want): + """Compare two floats by bit pattern, treating any two NaNs as equal. + + Bit patterns rather than ``==`` because ``-0.0 == 0.0``, and the sign of + a zero is exactly what a min/max tie-break decides. NaNs are exempted + because ``fmin`` may hand back either operand's NaN payload. + """ + if np.isnan(got) and np.isnan(want): + return True + return got.tobytes() == want.tobytes() + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +@pytest.mark.parametrize("np_dtype", [np.float64, np.float32]) +def test_udt_min_max_answer_what_the_builtin_dtype_answers(udt_op_path, np_dtype): + """``binary.min[udt]`` must give what ``binary.min[FP64]`` gives, bit for bit. + + An operator that means one thing on FP64 and another on a record of + FP64 is not one operator. SuiteSparse's ``GrB_MIN_FP64`` is C99 ``fmin``, + so that is what the UDT kernels have to be, and this compares them + directly against the built-in rather than against a convention chosen on + the Python side. The grid is every ordered pair drawn from NaN, both + infinities, both zeros and two ordinary values, so it covers a NaN on + either side, two NaNs, and a signed-zero tie either way round. + + The signed-zero tie itself is compared by value only. C99 leaves + ``fmin(-0.0, 0.0)`` unspecified and the built-in answers differently per + platform (left operand on macOS x86, right operand on Linux x86, IEEE + minNum on arm64 and Windows), so bit-for-bit agreement on that one pair + is not something any implementation can promise. Everything else, + including which zero a mixed zero/nonzero pair keeps, stays bit-exact. + + What this catches, in the two spellings it replaces: Python's builtin + ``min``, which the generated code reached through the exec namespace, + ordered NaN by position, and ``np.fmin`` under Numba gets the NaN rule + right but keeps the left operand on a signed-zero tie, so it drifts from + the JIT C kernel on ``min(0.0, -0.0)``. Both execution paths are checked + because SuiteSparse picks between them without telling anyone. + """ + nan, inf = float("nan"), float("inf") + values = [nan, inf, -inf, -0.0, 0.0, 1.5, -2.5] + xs = [x for x in values for _ in values] + ys = list(values) * len(values) + + udt = dtypes.register_anonymous( + np.dtype([("mmb_a", np_dtype)], align=True), f"_MinMaxBuiltin{np.dtype(np_dtype).name}" + ) + v, w = _udt_vectors(udt, xs, ys) + ref_v = Vector.from_dense(np.array(xs, dtype=np_dtype)) + ref_w = Vector.from_dense(np.array(ys, dtype=np_dtype)) + + for gb_op in (binary.min, binary.max): + expected = gb_op(ref_v & ref_w).new().to_dense() + result = gb_op(v & w).new() + for i, (x, y) in enumerate(zip(xs, ys, strict=True)): + got = result[i].new().value[0] + if x == 0 and y == 0 and np.signbit(x) != np.signbit(y): + # The one unspecified cell of the grid: either signed zero is + # a correct answer from either implementation, so only agree + # that both produced a zero. + msg = ( + f"{udt_op_path} {gb_op.name}({x}, {y}) on {udt.name}: " + f"got {got!r}, built-in {np.dtype(np_dtype).name} gives {expected[i]!r}" + ) + assert got == 0, msg + assert expected[i] == 0, msg + continue + assert _bitwise_eq(got, expected[i]), ( + f"{udt_op_path} {gb_op.name}({x}, {y}) on {udt.name}: " + f"got {got!r}, built-in {np.dtype(np_dtype).name} gives {expected[i]!r}" + ) + + # A NaN anywhere in the input must not change where a reduce lands. Under + # the Python-builtin semantics this same multiset reduced to 1.0 or to nan + # depending on which index the NaN sat at. + for data in ([1.0, 2.0, 3.0, nan], [nan, 1.0, 2.0, 3.0], [1.0, nan, 3.0, 2.0]): + (u,) = _udt_vectors(udt, data) + assert u.reduce(monoid.min[udt]).new().value[0] == 1.0, f"{udt_op_path} {data}" + assert u.reduce(monoid.max[udt]).new().value[0] == 3.0, f"{udt_op_path} {data}" + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_truediv_divides_in_floating_point(udt_op_path): + """``binary.truediv`` on integer fields must divide the way Python does. + + Regression: the JIT kernel emitted C ``/``, which is integer division. + ``10**18 / 3`` came out as 333333333333333333 under the JIT and + 333333333333333312 (float64, like numpy) through the cfunc, so the same + program gave different answers depending on whether a C compiler was + installed. + """ + udt = dtypes.register_anonymous( + np.dtype([("tdv_i", np.int64), ("tdv_j", np.int64)], align=True), "_TrueDivIntUDT" + ) + xs = [10**18, 10**18 + 1, 7, 22] + ys = [3, 3, 2, 7] + expected = (np.array(xs, np.int64) / np.array(ys, np.int64)).astype(np.int64) + assert expected[0] == 333333333333333312 # not 333333333333333333 + v, w = _udt_vectors(udt, xs, ys) + result = binary.truediv(v & w).new() + got = [result[i].new().value[0] for i in range(len(xs))] + assert got == list(expected), udt_op_path + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_floordiv_matches_numpy_on_floats(udt_op_path): + """``binary.floordiv`` on float fields is not ``floor(a / b)``. + + Regression: the JIT kernel computed ``floor(a / b)``, which rounds + differently from the remainder-based algorithm numpy and CPython use and + treats infinities as ordinary values. ``1.0 // 0.1`` came out as 10.0 + instead of 9.0, and ``inf // 2.0`` as ``inf`` instead of NaN, while the + cfunc agreed with numpy all along. + """ + nan = float("nan") + inf = float("inf") + # ``floor(a / b)`` disagrees with numpy on the first four pairs: inf and + # -inf where numpy gives NaN, 10.0 rather than 9.0 for 1.0 // 0.1, and + # -0.0 rather than -1.0 for -2.0 // inf. + xs = [inf, -inf, 1.0, -2.0, 2.0, -2.0, 0.0, nan, -7.0, 7.0, -0.0, 7.5] + ys = [2.0, 2.0, 0.1, inf, 0.0, 0.0, 0.0, 2.0, 2.0, -2.0, 4.0, 2.5] + for np_dtype, name in ((np.float64, "_FloorDivF64UDT"), (np.float32, "_FloorDivF32UDT")): + udt = dtypes.register_anonymous( + np.dtype([("fdv_a", np_dtype), ("fdv_b", np_dtype)], align=True), name + ) + with np.errstate(divide="ignore", invalid="ignore"): + expected = np.floor_divide(np.array(xs, np_dtype), np.array(ys, np_dtype)) + v, w = _udt_vectors(udt, xs, ys) + result = binary.floordiv(v & w).new() + got = np.array([result[i].new().value[0] for i in range(len(xs))], np_dtype) + np.testing.assert_array_equal(got, expected, err_msg=f"{udt_op_path} {np_dtype.__name__}") + # ``assert_array_equal`` reads -0.0 and 0.0 as equal, so the sign of a + # zero quotient needs its own assertion. It is the whole job of the + # ``copysign`` branch the JIT kernel emits for an exact-zero result. + np.testing.assert_array_equal( + np.signbit(got), + np.signbit(expected), + err_msg=f"{udt_op_path} {np_dtype.__name__} sign of zero", + ) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_integer_division_by_zero_is_defined(udt_op_path): + """Integer division by zero must return a value rather than trap. + + The JIT kernel divided in integers, where a zero divisor is undefined + behaviour: on x86-64 ``idiv`` raises #DE, which is SIGFPE and process + death rather than an exception. AArch64's ``sdiv`` returns 0 and does not + trap, so this cannot be exhibited on an arm64 machine. The same trap + fires on ``INT_MIN / -1``, whose quotient is not representable. + + The values below are a choice, not a discovery. A quotient that doesn't + fit the field is an undefined conversion in both C and LLVM, and the two + answered differently per field width, so both paths now rule those cases + out: a zero divisor gives 0, as ``np.floor_divide`` does, and + ``INT_MIN // -1`` wraps to ``INT_MIN``, as numpy does. + """ + signed = dtypes.register_anonymous( + np.dtype([("dvz_a", np.int32), ("dvz_b", np.int8)], align=True), "_DivZeroSignedUDT" + ) + unsigned = dtypes.register_anonymous( + np.dtype([("dvz_c", np.uint32), ("dvz_d", np.uint64)], align=True), "_DivZeroUnsignedUDT" + ) + v, w = _udt_vectors(signed, [7, -7, 100, -128], [0, 0, 3, -1]) + for gb_op in (binary.truediv, binary.floordiv): + result = gb_op(v & w).new() + got = [result[i].new().value[0] for i in range(4)] + assert got[:2] == [0, 0], f"{udt_op_path} {gb_op.name}" + assert got[2] == 33, f"{udt_op_path} {gb_op.name}" # 100 / 3 truncates either way + # ``-128 // -1`` is the second trapping case; numpy wraps it to INT8_MIN. + result = binary.floordiv(v & w).new() + assert result[3].new().value[1] == np.iinfo(np.int8).min, udt_op_path + + v, w = _udt_vectors(unsigned, [7, 9, 100, 5], [0, 0, 3, 2]) + for gb_op in (binary.truediv, binary.floordiv): + result = gb_op(v & w).new() + got = [result[i].new().value[0] for i in range(4)] + assert got == [0, 0, 33, 2], f"{udt_op_path} {gb_op.name}" + + # Floor division still floors for signed operands of mixed sign. + v, w = _udt_vectors(signed, [-7, 7, -9, 11], [2, -2, 2, 3]) + result = binary.floordiv(v & w).new() + got = [result[i].new().value[0] for i in range(4)] + assert got == [-4, -4, -5, 3], udt_op_path + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.skipif("not dtypes._supports_complex") +def test_udt_complex_truediv_by_zero(udt_op_path): + """``binary.truediv`` on a complex field survives a zero divisor. + + Numba's complex division raises ``ZeroDivisionError`` unconditionally, + outside the error model's control, so the cfunc left the element unwritten + while the JIT kernel returned numpy's infinities. Reading back the + abandoned element gave uninitialized memory, or the previous element's + answer, either of which looks like a plausible value. + """ + udt = dtypes.register_anonymous( + np.dtype([("cxz_a", np.complex128)], align=True), "_ComplexDivZeroUDT" + ) + xs = [3 + 4j, 0j, 1 + 1j, 2 - 2j] + ys = [0j, 0j, 2 + 0j, 0j] + v, w = _udt_vectors(udt, xs, ys) + result = binary.truediv(v & w).new() + got = np.array([result[i].new().value[0] for i in range(len(xs))]) + with np.errstate(divide="ignore", invalid="ignore"): + expected = np.array(xs) / np.array(ys) + np.testing.assert_array_equal(got, expected, err_msg=udt_op_path) + + +@pytest.mark.skipif("not supports_udfs") +# 136-byte UDT, which SS < 9 rejects; see test_udt_large_array. +@pytest.mark.skipif( + "ss_version_major < 9", + reason="SuiteSparse < 9 rejects a 136-byte UDT on builds without VLA support", +) +def test_udt_float_truediv_by_zero_is_infinite(udt_op_path): + """A zero divisor on a float field gives numpy's infinity, not a lost element. + + Unlike the integer case, nothing here is guarded: the generated code + divides and lets IEEE produce the infinity. That only holds because the + generated wrapper is compiled under Numba's numpy error model, which + nothing else in the suite pins down. + """ + udt = dtypes.register_anonymous(np.dtype((np.float64, (17,))), "_FloatDivZeroArr17") + xs = [1.0, -1.0, 0.0, 6.0] + ys = [0.0, 0.0, 0.0, 3.0] + v, w = _udt_vectors(udt, xs, ys) + result = binary.truediv(v & w).new() + got = np.array([result[i].new().value[0] for i in range(len(xs))]) + with np.errstate(divide="ignore", invalid="ignore"): + expected = np.array(xs) / np.array(ys) + np.testing.assert_array_equal(got, expected, err_msg=udt_op_path) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_array_ops_match_record_ops(udt_op_path): + """The array-UDT codegen carries the same division and NaN fixes as records. + + Records and flat arrays go through separate branches in both the Numba + and the JIT C generators, so each fix has to land in both. + """ + nan = float("nan") + inf = float("inf") + float_udt = dtypes.register_anonymous(np.dtype((np.float64, (13,))), "_ArrOpsF64") + xs = [inf, 1.0, -2.0, nan, -7.0, 2.0] + ys = [2.0, 0.1, inf, 2.0, 2.0, 1.0] + v, w = _udt_vectors(float_udt, xs, ys) + # The reference computations touch inf and nan, and numpy raises the FP + # invalid flag for them on some platforms (Linux and Windows, via fmod) + # but not others; pyproject promotes the RuntimeWarning to an error. + with np.errstate(divide="ignore", invalid="ignore"): + expected_floordiv = np.floor_divide(np.array(xs), np.array(ys)) + expected_min = np.fmin(np.array(xs), np.array(ys)) + np.testing.assert_array_equal( + [binary.floordiv(v & w).new()[i].new().value[0] for i in range(len(xs))], + expected_floordiv, + err_msg=udt_op_path, + ) + # ``np.fmin``, not ``np.minimum``: ``binary.min`` is SuiteSparse's + # ``GrB_MIN_FP64``, which ignores a NaN operand rather than propagating it. + np.testing.assert_array_equal( + [binary.min(v & w).new()[i].new().value[0] for i in range(len(xs))], + expected_min, + err_msg=udt_op_path, + ) + + int_udt = dtypes.register_anonymous(np.dtype((np.int64, (6,))), "_ArrOpsI64") + ixs = [10**18, 7, -7, 100, -9, 5] + iys = [3, 0, 0, 3, 2, 2] + v, w = _udt_vectors(int_udt, ixs, iys) + result = binary.truediv(v & w).new() + got = [result[i].new().value[0] for i in range(len(ixs))] + assert got == [333333333333333312, 0, 0, 33, -4, 2], udt_op_path + result = binary.floordiv(v & w).new() + got = [result[i].new().value[0] for i in range(len(ixs))] + assert got == [333333333333333333, 0, 0, 33, -5, 2], udt_op_path + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_multidim_array_ops_match_numpy(udt_op_path): + """Built-in ops on a multi-dimensional array UDT agree with numpy on both paths. + + The JIT C typedef flattens any rank to ``double v [N]`` and the Numba + wrapper walks the same flat run, so a 2-D UDT covers codegen that the 1-D + cases reach only by accident of both being contiguous. + """ + udt = dtypes.register_anonymous(np.dtype((np.float64, (3, 2))), "_ArrOps2D") + xs = [1.0, -7.0, float("inf"), 2.0] + ys = [0.1, 2.0, 2.0, 0.0] + v, w = _udt_vectors(udt, xs, ys) + for gb_op, reference in ( + (binary.floordiv, np.floor_divide), + (binary.truediv, np.true_divide), + # ``fmin`` rather than ``minimum``: these inputs carry no NaN, so the + # two agree here, but ``binary.min`` is the NaN-ignoring one. + (binary.min, np.fmin), + ): + result = gb_op(v & w).new() + element = result[0].new().value + assert element.shape == (3, 2) + got = np.array([result[i].new().value[1, 1] for i in range(len(xs))]) + with np.errstate(divide="ignore", invalid="ignore"): + expected = reference(np.array(xs), np.array(ys)) + np.testing.assert_array_equal(got, expected, err_msg=f"{udt_op_path} {gb_op.name}") + + @pytest.mark.skipif("not supports_udfs") @pytest.mark.slow def test_udt_tuple_return_binaryop(record_udt): @@ -1859,6 +2366,405 @@ def test_udt_auto_monoid(): assert any_res in ((7, 8.0), (11, 12.0)) +@pytest.mark.skipif("not supports_udfs") +def test_udt_array_wrapper_stays_within_element(): + """The array-UDT cfunc wrapper writes the element payload and nothing more. + + Numba represents a ``NestedArray`` *value* as an array descriptor (data + pointer, shape, strides, ...), so the wrapper used to store that descriptor + where GraphBLAS expected only the elements: 56 bytes into the 48 an + ``FP64[6]`` element gets. GraphBLAS owns the buffer it hands the cfunc, so + the extra bytes land on memory the library allocated for something else. + + Pinned at the codegen level rather than end-to-end, and deliberately so: + nothing in the suite was ever observed to fault on the old codegen, so an + end-to-end test would not catch a regression. Measured directly instead, + driving the wrapper through ctypes over a guard-filled buffer: shape + ``(6,)`` wrote 8 bytes past a 48-byte element, shape ``(2, 3)`` wrote 24. + """ + import ctypes + + import numba + + from graphblas.core.operator.base import _get_udt_wrapper + + size = 6 + udt = dtypes.register_anonymous(np.dtype((np.float64, (size,))), "_ArrWrapPin") + + @numba.njit + def _second(x, y): # pragma: no cover (numba) + return y + + wrapper, wrapper_sig = _get_udt_wrapper(_second, udt, udt, udt) + cfunc = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")(wrapper) + call = ctypes.CFUNCTYPE(None, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p)(cfunc.address) + + guard = -1.0 + xvals = [float(i) for i in range(size)] + yvals = [100.0 + i for i in range(size)] + z = (ctypes.c_double * (4 * size))(*([guard] * (4 * size))) + x = (ctypes.c_double * size)(*xvals) + y = (ctypes.c_double * size)(*yvals) + call(ctypes.byref(z), ctypes.byref(x), ctypes.byref(y)) + + assert list(z)[:size] == yvals + assert all(val == guard for val in list(z)[size:]), f"wrote past the UDT element: {list(z)}" + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_array_any_wrapper_stays_within_element(): + """``binary.any`` on an array UDT must write one element, not Numba's array descriptor. + + ``any``, ``first``, and ``second`` are not in ``_BUILTIN_UDT_BINARY_OPS``, + so they compile through the generic ``_numba_func`` branch of + ``BinaryOp._compile_udt``. The wrapper there used to load and store the + operand as a ``NestedArray`` *value*, which Numba models as its full array + descriptor (meminfo, parent, nitems, itemsize, data, shape, strides): 56 + bytes on 64-bit for a 1-D element, regardless of payload size. SuiteSparse's + generic reduce keeps a UDT accumulator in a stack array sized to the element + (32 bytes here), so each fold overflowed it by 24 bytes; depending on the + build that clobbered a spilled pointer (segfault or SIGBUS) or silently + produced a wrong answer. + + Drive the compiled wrapper directly, over heap buffers with slack, so a + regression trips an assert instead of corrupting a stack frame. The two + sentinels must differ: the descriptor load/store is a byte-preserving copy + of the source element plus its trailing bytes, so if ``y``'s slack held the + same sentinel as ``z``'s guard, the overflow would rewrite ``z``'s guard + bytes with identical values and the check would be blind to it. + """ + import ctypes + + import numba + + from graphblas.core.operator.base import _get_udt_wrapper + + # ``register_anonymous`` caches per np.dtype, so this may return the same + # DataType as other float64[4] tests, renamed. That is fine here: the + # wrapper below is compiled fresh and nothing asserts on cached JIT state. + udt = dtypes.register_anonymous(np.dtype((np.float64, (4,))), "_AnyOverflowArr") + + # Mirror the generic ``_numba_func`` branch of ``BinaryOp._compile_udt``. + numba_func = binary.any._numba_func + sig = (udt.numba_type, udt.numba_type) + numba_func.compile(sig) + numba_ret_type = numba_func.overloads[sig].signature.return_type + wrapper, wrapper_sig = _get_udt_wrapper( + numba_func, udt, udt, udt, numba_ret_type=numba_ret_type + ) + cfunc = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")(wrapper) + call = ctypes.CFUNCTYPE(None, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p)(cfunc.address) + + itemsize = udt.np_type.itemsize + slack = 128 # the descriptor overran by 24 bytes; leave generous headroom + z = np.full(itemsize + slack, 0xAB, dtype=np.uint8) + x = np.full(itemsize + slack, 0xCD, dtype=np.uint8) + y = np.full(itemsize + slack, 0xCD, dtype=np.uint8) + xvals = np.array([1.0, 2.0, 3.0, 4.0]) + yvals = np.array([10.0, 20.0, 30.0, 40.0]) + x[:itemsize] = xvals.view(np.uint8) + y[:itemsize] = yvals.view(np.uint8) + call(z.ctypes.data, x.ctypes.data, y.ctypes.data) + + # ``any`` uses ``_second`` semantics, so the payload must be ``y``'s. + np.testing.assert_array_equal(z[:itemsize].view(np.float64), yvals) + overrun = np.flatnonzero(z[itemsize:] != 0xAB) + assert overrun.size == 0, f"wrote {overrun.size} bytes past the element at offsets {overrun}" + + # Public-path smoke: the reduce whose stack accumulator the old wrapper + # overflowed. Kept after the byte-level checks so a regression fails the + # assert above instead of reaching code that may crash the process. + v = Vector(udt, size=3) + rows = [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0], [9.0, 10.0, 11.0, 12.0]] + for i, row in enumerate(rows): + v[i] = row + res = v.reduce(monoid.any).new() + assert any(np.array_equal(res.value, row) for row in rows) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_array_udf_returns_new_array(): + """An array-UDT UDF may build its result instead of returning an operand. + + The wrapper hands the UDF a numpy view of each element in the UDT's + declared shape, so ordinary array expressions work and the result is copied + back element-wise. + """ + udt = dtypes.register_anonymous(np.dtype((np.float64, (8,))), "_ArrRetUDT") + v = Vector(udt, size=2) + v[0] = np.arange(8.0) + v[1] = np.arange(8.0, 16.0) + w = Vector(udt, size=2) + w[0] = np.full(8, 100.0) + w[1] = np.full(8, 200.0) + + def _add(x, y): + return x + y # pragma: no cover (numba) + + add_op = BinaryOp.register_anonymous(_add, "_arr_ret_add", is_udt=True) + result = add_op(v & w).new() + np.testing.assert_array_equal(result[0].new().value, np.arange(8.0) + 100.0) + np.testing.assert_array_equal(result[1].new().value, np.arange(8.0, 16.0) + 200.0) + + def _double(x): + return x * 2 # pragma: no cover (numba) + + double_op = UnaryOp.register_anonymous(_double, "_arr_ret_double", is_udt=True) + np.testing.assert_array_equal(double_op(v).new()[0].new().value, np.arange(8.0) * 2) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_multidim_array_keeps_shape_in_udf(): + """A multi-dimensional array UDT reaches the UDF in its declared shape. + + The wrapper builds the operand with ``numba.carray(ptr, shape)``, so 2-D + indexing and ``.shape`` work. Addressing it as a flat run of elements + would still be memory-safe but would silently drop that metadata. + """ + udt2d = dtypes.register_anonymous(np.dtype((np.float64, (2, 3))), "_ArrRet2D") + m = Vector(udt2d, size=1) + m[0] = np.arange(6.0).reshape(2, 3) + n = Vector(udt2d, size=1) + n[0] = np.full((2, 3), 10.0) + + def _add_corner(x, y): + # Fails to compile unless `x` really is 2-D with shape metadata. + return x + y[0, 0] + x.shape[1] # pragma: no cover (numba) + + add_2d = BinaryOp.register_anonymous(_add_corner, "_arr_ret_add_2d", is_udt=True) + np.testing.assert_array_equal( + add_2d(m & n).new()[0].new().value, np.arange(6.0).reshape(2, 3) + 10.0 + 3 + ) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_array_udf_shape_errors(): + """Array-UDT UDFs that can't fill the element are rejected at registration. + + Numba's ``Array`` type records ``ndim`` but not extents, so neither case + below is a type error. Both used to reach the cfunc, where the shape + mismatch raises in a context that swallows the exception, handing the + caller an uninitialized element and no error. + """ + # Shapes unique to this test: ``register_anonymous`` caches by dtype and + # freezes the JIT C name at first registration, so sharing a shape with + # another test makes both order-dependent. + udt9 = dtypes.register_anonymous(np.dtype((np.float64, (9,))), "_ShapeErr9") + udt10 = dtypes.register_anonymous(np.dtype((np.float64, (10,))), "_ShapeErr10") + + def _truncate(x): # pragma: no cover (numba) + return x[:2] + + op = UnaryOp.register_anonymous(_truncate, "_shape_err_trunc", is_udt=True) + with pytest.raises(UdfParseError, match=r"shape \(2,\) when run on sample values"): + op[udt9] + + # Two array UDTs sharing a base dtype and rank are indistinguishable once a + # UDF builds its result, so refuse to guess which one it meant. + def _built(x, y): # pragma: no cover (numba) + return y + 0.0 + + op2 = BinaryOp.register_anonymous(_built, "_shape_err_ambiguous", is_udt=True) + with pytest.raises(UdfParseError, match="matches more than one input array UDT"): + op2[udt9, udt10] + + # Ambiguity is decided on the UDTs, not their Numba shapes: these two are + # separate DataTypes with separate GraphBLAS handles, but Numba collapses + # the layered dtype to the flat one's ``nestedarray(float64, (2, 3))``. + flat = dtypes.register_anonymous(np.dtype((np.float64, (3, 4))), "_ShapeErrFlat") + layered = dtypes.register_anonymous( + np.dtype((np.dtype((np.float64, (4,))), (3,))), "_ShapeErrLayered" + ) + assert flat.numba_type == layered.numba_type + with pytest.raises(UdfParseError, match="matches more than one input array UDT"): + op2[flat, layered] + assert op2[flat, flat].return_type is flat # a same-type pair is not ambiguous + + # An array UDF whose result matches no input names the mismatch rather + # than telling the user to return an array, which is what they did. + def _recast(x): # pragma: no cover (numba) + return x.astype(np.float32) + + op3 = UnaryOp.register_anonymous(_recast, "_shape_err_recast", is_udt=True) + with pytest.raises(UdfParseError, match="matches no input array UDT"): + op3[udt9] + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_record_array_leaf_shape_errors(): + """A record UDF that under-fills an array-typed leaf is rejected at registration. + + The wrapper slice-assigns array leaves, so a short return raises inside + the cfunc and abandons the write part-way: leaves after it keep whatever + SuiteSparse had in the buffer, scalar leaves included. + """ + spec = np.dtype([("rl_vec", np.float64, (3,)), ("rl_tag", np.int64)], align=True) + udt = dtypes.register_anonymous(spec, "_RecLeafShape") + + def _short(x, y): # pragma: no cover (numba) + return (x["rl_vec"][:2], x["rl_tag"]) + + op = BinaryOp.register_anonymous(_short, "_rec_leaf_short", is_udt=True) + with pytest.raises(UdfParseError, match=r"shape \(2,\) for field .* holds \(3,\)"): + op[udt] + + def _full(x, y): # pragma: no cover (numba) + return (x["rl_vec"] + y["rl_vec"], x["rl_tag"] + y["rl_tag"]) + + op = BinaryOp.register_anonymous(_full, "_rec_leaf_full", is_udt=True) + v = Vector(udt, size=1) + v[0] = ([1.0, 2.0, 3.0], 7) + w = Vector(udt, size=1) + w[0] = ([4.0, 5.0, 6.0], 8) + got = v.ewise_mult(w, op).new()[0].new().value + np.testing.assert_array_equal(got["rl_vec"], [5.0, 7.0, 9.0]) + assert got["rl_tag"] == 15 + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_array_udf_broadcast_return(): + """A return that broadcasts to the element fills it, and is not rejected. + + The wrapper slice-assigns and numpy broadcasts on assignment, so a ``(1,)`` + return legitimately fills every slot of a ``(6,)`` element. Requiring an + exact shape would refuse this, which works. + """ + udt6 = dtypes.register_anonymous(np.dtype((np.float64, (6,))), "_BCast6") + + def _fill(x): # pragma: no cover (numba) + return x[:1] + 10.0 + + op1 = UnaryOp.register_anonymous(_fill, "_bcast_fill", is_udt=True) + assert op1[udt6].return_type is udt6 + v = Vector(udt6, size=1) + v[0] = np.arange(1.0, 7.0) + np.testing.assert_array_equal(v.apply(op1).new()[0].new().value, [11.0] * 6) + + # A row broadcast across a 2-D element: the same rule one rank up. + udt42 = dtypes.register_anonymous(np.dtype((np.float64, (4, 2))), "_BCast42") + + def _fill_rows(x): # pragma: no cover (numba) + return x[:1, :] + 100.0 + + op2 = UnaryOp.register_anonymous(_fill_rows, "_bcast_fill_rows", is_udt=True) + assert op2[udt42].return_type is udt42 + v2 = Vector(udt42, size=1) + v2[0] = np.arange(8.0).reshape(4, 2) + np.testing.assert_array_equal( + v2.apply(op2).new()[0].new().value, np.tile([100.0, 101.0], (4, 1)) + ) + + # The other side of the boundary: (2,) does not broadcast to (6,), Numba's + # slice-assign raises on it, and it stays rejected. + def _short(x): # pragma: no cover (numba) + return x[:2] + 10.0 + + op3 = UnaryOp.register_anonymous(_short, "_bcast_short", is_udt=True) + with pytest.raises(UdfParseError, match=r"shape \(2,\) when run on sample values"): + op3[udt6] + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_record_leaf_broadcast_return(): + """A broadcastable array leaf fills its field, and later leaves still land. + + Same boundary as the array case, and + ``test_udt_record_array_leaf_shape_errors`` holds the rejecting side. The + scalar leaf is worth asserting because a leaf that raises in the cfunc + abandons the write, leaving every leaf after it as SuiteSparse had it. + """ + spec = np.dtype([("bc_vec", np.float64, (11,)), ("bc_tag", np.int64)], align=True) + udt = dtypes.register_anonymous(spec, "_RecLeafBCast") + + def _fill_leaf(x, y): # pragma: no cover (numba) + return (x["bc_vec"][:1] + y["bc_vec"][:1], x["bc_tag"] + y["bc_tag"]) + + op1 = BinaryOp.register_anonymous(_fill_leaf, "_rec_leaf_bcast", is_udt=True) + v = Vector(udt, size=1) + v[0] = (np.arange(11.0), 7) + w = Vector(udt, size=1) + w[0] = (np.arange(11.0) + 1.0, 8) + got = v.ewise_mult(w, op1).new()[0].new().value + np.testing.assert_array_equal(got["bc_vec"], [1.0] * 11) + assert got["bc_tag"] == 15 + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_broadcast_matches_numba_slice_assign(): + """The shape check accepts exactly what the wrapper's slice-assign accepts. + + The check turns a silent cfunc failure into a registration error, so a + shape it rejects that Numba would have assigned is a false rejection, and + one it accepts that Numba raises on is the failure it exists to catch. Pin + both directions against Numba itself, including the two ranks where + broadcasting alone gives the wrong answer: ``(1, 6)`` fills a ``(6,)`` + destination because assignment drops leading ones, ``(6, 1)`` does not. + """ + import numba + + from graphblas.core.operator.base import _fits_by_broadcast + + @numba.njit + def _assign(z, src): # pragma: no cover (numba) + z[:] = src + + for dst, src in [ + ((6,), ()), + ((6,), (1,)), + ((6,), (6,)), + ((6,), (2,)), + ((6,), (12,)), + ((6,), (1, 6)), + ((6,), (6, 1)), + ((2, 3), (1, 3)), + ((2, 3), (2, 1)), + ((2, 3), (1, 1)), + ((2, 3), (3,)), + ((2, 3), (2, 3)), + ((2, 3), (6,)), + ((2, 3), (3, 2)), + ]: + try: + _assign(np.zeros(dst), np.ones(src)) + except ValueError: + numba_assigns = False + else: + numba_assigns = True + assert _fits_by_broadcast(src, dst) is numba_assigns, (src, dst) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_record_array_field_roundtrip(): + """A record UDT with an array field writes exactly that field's extent. + + Numba's record-field setitem copies the *destination* extent whatever the + source's length, so a short return used to read past the end of the source + array. The wrapper slice-assigns array leaves to make that a shape error. + """ + spec = np.dtype([("count", np.int64), ("vec", np.float64, (3,))], align=True) + udt = dtypes.register_anonymous(spec, "_RecArrField") + v = Vector(udt, size=2) + v[0] = (1, [1.0, 2.0, 3.0]) + v[1] = (2, [4.0, 5.0, 6.0]) + + def _combine(x, y): # pragma: no cover (numba) + return (x["count"] + y["count"], x["vec"] + y["vec"]) + + op = BinaryOp.register_anonymous(_combine, "_rec_arr_field", is_udt=True) + result = v.ewise_mult(v, op).new() + assert result[0].new().value["count"] == 2 + np.testing.assert_array_equal(result[0].new().value["vec"], [2.0, 4.0, 6.0]) + np.testing.assert_array_equal(result[1].new().value["vec"], [8.0, 10.0, 12.0]) + + @pytest.mark.skipif("not supports_udfs") @pytest.mark.slow def test_udt_auto_semiring(): @@ -2385,6 +3291,42 @@ def test_udt_eq_ne_rejects_incompatible_pairs(): binary.eq(v_uv & v_arr).new() +@pytest.mark.skipif("not supports_udfs") +# SS < 9 has no GrB_NAME setter, so registration falls back to storing the +# numpy repr in the type name and warns when it does not fit in 128 chars. +# _NestDeep's repr is 142; how it serializes is not what the test is about. +@pytest.mark.filterwarnings("ignore:UDT repr is too large") +def test_udt_record_nesting_mismatch_is_a_keyerror(): + """Records sharing field names but not nesting depth are rejected as a KeyError. + + ``_check_udt_pair`` matched on top-level names only, but the codegen pairs + operands leaf by leaf, and a field that is a sub-record on one side and a + scalar on the other contributes a different number of leaves. Without the + guard the pair reaches Numba, whose typing failure arrives as a + ``UdfParseError``: a compile error reported for what is really the same + shape disagreement its sibling checks raise ``KeyError`` for. + """ + flat = dtypes.register_anonymous( + np.dtype([("nst_a", np.float64), ("nst_b", np.float64)], align=True), "_NestFlat" + ) + nested = dtypes.register_anonymous( + np.dtype( + [ + ("nst_a", np.dtype([("nst_n1", np.float64), ("nst_n2", np.float64)])), + ("nst_b", np.float64), + ], + align=True, + ), + "_NestDeep", + ) + v = Vector(flat, size=1) + v[0] = (1.0, 2.0) + w = Vector(nested, size=1) + w[0] = ((3.0, 4.0), 5.0) + with pytest.raises(KeyError, match="same number of leaf fields"): + v.ewise_mult(w, binary.plus).new() + + @pytest.mark.skipif("not supports_udfs") @pytest.mark.slow def test_udt_aggregators(): @@ -2796,3 +3738,434 @@ def test_compile_codegen_helper(): assert "Source:" in msg assert bad_src in msg assert isinstance(exc_info.value.__cause__, SyntaxError) + + +def test_operator_namespace_typo_suggestions(): + # A typo in an operator namespace should suggest close matches (via difflib), + # drawn from __dir__() so lazily-registered operators are offered without + # forcing them to build. + with pytest.raises(AttributeError, match="has no attribute 'pluss'.*Did you mean 'plus'"): + binary.pluss + with pytest.raises(AttributeError, match="Did you mean 'plus'"): + monoid.pluss + with pytest.raises(AttributeError, match="plus_times"): + semiring.plus_time + with pytest.raises(AttributeError, match="Did you mean 'sum'"): + agg.summ + with pytest.raises(AttributeError, match="Did you mean"): + unary.expp + with pytest.raises(AttributeError, match="rowindex"): + indexunary.rowindexx + with pytest.raises(AttributeError, match="triu"): + select.triu_typo + with pytest.raises(AttributeError, match="Did you mean 'plus'"): + op.pluss + + # No close match -> plain message, no suggestion appended + with pytest.raises(AttributeError) as exc_info: + binary.zzzzzz + assert "has no attribute 'zzzzzz'" in str(exc_info.value) + assert "Did you mean" not in str(exc_info.value) + + # Building suggestions must not force lazy operators to compile + before = set(binary._delayed) + with pytest.raises(AttributeError): + binary.pluss + assert set(binary._delayed) == before + + +# Touching an operator namespace must not compile any lazily-registered UDF. +# Run in a subprocess: BinaryOp._initialize runs once per process, so by the +# time any test executes, the import-time behavior under test is long past. +_LAZY_UDF_PROBE = """ +import graphblas as gb + +gb.binary.plus # forces BinaryOp._initialize + +lazy_udfs = ("floordiv", "rfloordiv", "absfirst", "abssecond", "rpow") +print("package: " + gb.__file__) +print("still lazy: " + " ".join(n for n in lazy_udfs if n in gb.binary._delayed)) +""" + + +@pytest.mark.skipif("not supports_udfs") +def test_initialize_does_not_build_lazy_udfs(): + repo_root = Path(__file__).resolve().parents[2] + env = dict(os.environ, PYTHONPATH=str(repo_root)) + result = subprocess.run( + [sys.executable, "-c", _LAZY_UDF_PROBE], + capture_output=True, + text=True, + check=False, + env=env, + ) + report = f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + assert result.returncode == 0, report + + lines = dict(line.split(": ", 1) for line in result.stdout.splitlines() if ": " in line) + # Assert the child probed this tree before trusting what it reports about it. + assert lines.get("package") == str(repo_root / "graphblas" / "__init__.py"), report + + still_lazy = lines.get("still lazy", "").split() + assert still_lazy == ["floordiv", "rfloordiv", "absfirst", "abssecond", "rpow"], report + + +@pytest.mark.skipif("not supports_udfs") +def test_builtin_udfs_are_disk_cached(): + # The built-in UDF binops are module-level functions, so numba can persist + # their compilation across processes. Anything a user registers cannot be: + # numba keys a cache entry on a stable on-disk source, which a lambda or an + # interactively defined function does not have. + from numba.core.caching import NullCache + + for name in ["floordiv", "rfloordiv", "absfirst", "abssecond", "rpow"]: + numba_func = getattr(binary, name)._numba_func + assert not isinstance(numba_func._cache, NullCache), name + + def _uncached_probe(x, y): + return x + y + + user_op = BinaryOp.register_anonymous(_uncached_probe) + assert isinstance(user_op._numba_func._cache, NullCache) + + +# The built-in UDF binops advertise every dtype in ``.types`` up front but +# compile none of them until asked. Run in a subprocess: any earlier test may +# already have materialized them in this process. +_DEFERRED_BUILD_PROBE = """ +import graphblas as gb + +op = gb.binary.floordiv +print("package: " + gb.__file__) +print("types: %d" % len(op.types)) +print("compiled before: %d" % len(op._typed_ops)) +op[gb.dtypes.INT64] +print("compiled after: %d" % len(op._typed_ops)) +""" + + +@pytest.mark.skipif("not supports_udfs") +def test_builtin_udf_types_precede_compilation(): + repo_root = Path(__file__).resolve().parents[2] + env = dict(os.environ, PYTHONPATH=str(repo_root)) + result = subprocess.run( + [sys.executable, "-c", _DEFERRED_BUILD_PROBE], + capture_output=True, + text=True, + check=False, + env=env, + ) + report = f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + assert result.returncode == 0, report + + lines = dict(line.split(": ", 1) for line in result.stdout.splitlines() if ": " in line) + assert lines.get("package") == str(repo_root / "graphblas" / "__init__.py"), report + + assert int(lines["types"]) == len(binary.floordiv.types), report + assert int(lines["compiled before"]) == 0, report + assert int(lines["compiled after"]) == 1, report + + +_DEFERRED_COMMUTES_PROBE = """ +import graphblas as gb + +print("package: " + gb.__file__) +ct = gb.binary.floordiv[gb.dtypes.INT64].commutes_to +print("floordiv_ok: " + str(ct is gb.binary.rfloordiv[gb.dtypes.INT64])) +ct = gb.binary.absfirst[gb.dtypes.INT64].commutes_to +print("absfirst_ok: " + str(ct is gb.binary.abssecond[gb.dtypes.INT64])) +""" + + +@pytest.mark.skipif("not supports_udfs") +def test_deferred_commutes_to(): + # A deferred partner op must still answer commutes_to. The membership + # test consults .types, not ._typed_ops: with the latter, a fresh process + # answered None for floordiv[INT64].commutes_to until rfloordiv happened + # to be compiled, so the answer depended on access order. Subprocess for + # the same reason as test_initialize_does_not_build_lazy_udfs: in this + # process the partners may already be built. + repo_root = Path(__file__).resolve().parents[2] + env = dict(os.environ, PYTHONPATH=str(repo_root)) + result = subprocess.run( + [sys.executable, "-c", _DEFERRED_COMMUTES_PROBE], + capture_output=True, + text=True, + check=False, + env=env, + ) + report = f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + assert result.returncode == 0, report + lines = dict(line.split(": ", 1) for line in result.stdout.splitlines() if ": " in line) + assert lines.get("package") == str(repo_root / "graphblas" / "__init__.py"), report + assert lines.get("floordiv_ok") == "True", report + assert lines.get("absfirst_ok") == "True", report + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_ret_dtype_names_an_output_udt(): + """``ret_dtype`` names an output UDT that is not one of the operands. + + Without it the return type is inferred from what the UDF builds, and the + only names in scope are the input dtypes. That makes a rank-reducing op + such as FP64[9] -> FP64[3] unreachable: the inferred type is the input's, + and the shape check then rejects the shorter array the UDF returns. + """ + # Shapes unique to this test; see the note in test_udt_array_udf_shape_errors. + nine = dtypes.register_anonymous(np.dtype((np.float64, (9,))), "_RetD9") + three = dtypes.register_anonymous(np.dtype((np.float64, (3,))), "_RetD3") + + def _first_three(x): # pragma: no cover (numba) + return x[:3] + + without = UnaryOp.register_anonymous(_first_three, "_ret_dtype_without", is_udt=True) + with pytest.raises(UdfParseError, match=r"shape \(3,\) when run on sample values"): + without[nine] + + op_ = UnaryOp.register_anonymous(_first_three, "_ret_dtype_with", is_udt=True, ret_dtype=three) + assert op_[nine].return_type is three + + v = Vector(nine, size=2) + v[0] = np.arange(9.0) + w = op_(v).new() + assert w.dtype is three + np.testing.assert_array_equal(w[0].new().value, np.arange(3.0)) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_ret_dtype_binary_and_record(): + """``ret_dtype`` on a binary op, including a type that is neither operand.""" + a4 = dtypes.register_anonymous(np.dtype((np.float64, (4,))), "_RetDBinA4") + b6 = dtypes.register_anonymous(np.dtype((np.float64, (6,))), "_RetDBinB6") + out2 = dtypes.register_anonymous(np.dtype((np.float64, (2,))), "_RetDBinOut2") + + def _head_two(x, y): # pragma: no cover (numba) + return x[:2] + y[:2] + + op_ = BinaryOp.register_anonymous(_head_two, "_ret_dtype_bin", is_udt=True, ret_dtype=out2) + assert op_[a4, b6].return_type is out2 + + # ret_dtype equal to an operand's dtype agrees with what inference picks. + def _plus(x, y): # pragma: no cover (numba) + return x + y + + same = BinaryOp.register_anonymous(_plus, "_ret_dtype_same", is_udt=True, ret_dtype=a4) + inferred = BinaryOp.register_anonymous(_plus, "_ret_dtype_inferred", is_udt=True) + assert same[a4, a4].return_type is inferred[a4, a4].return_type is a4 + + # A record UDT output that appears in no operand. + rec = dtypes.register_anonymous( + np.dtype([("lo", np.float64), ("hi", np.float64)], align=True), "_RetDRec" + ) + + def _bounds(x, y): # pragma: no cover (numba) + return (min(x, y), max(x, y)) + + recop = BinaryOp.register_anonymous(_bounds, "_ret_dtype_rec", is_udt=True, ret_dtype=rec) + assert recop[FP64, FP64].return_type is rec + + w = Vector(FP64, size=2) + w[0] = 3.0 + u = Vector(FP64, size=2) + u[0] = 1.0 + res = recop(w & u).new() + assert res.dtype is rec + assert tuple(res[0].new().value.tolist()) == (1.0, 3.0) + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_ret_dtype_errors(): + """``ret_dtype`` is rejected outside the UDT path and for unrecognized dtypes.""" + with pytest.raises(ValueError, match="not a recognized dtype"): + UnaryOp.register_anonymous(lambda x: x, "_ret_dtype_junk", is_udt=True, ret_dtype="NOPE") + with pytest.raises(ValueError, match="not a recognized dtype"): + UnaryOp.register_anonymous(lambda x: x, "_ret_dtype_junk2", is_udt=True, ret_dtype=object()) + + # The builtin path derives a return type per input dtype, so one fixed + # dtype cannot describe it. + with pytest.raises(ValueError, match="ret_dtype requires is_udt=True"): + UnaryOp.register_anonymous(lambda x: x, "_ret_dtype_builtin", ret_dtype=FP64) + with pytest.raises(ValueError, match="ret_dtype requires is_udt=True"): + BinaryOp.register_anonymous(lambda x, y: x + y, "_ret_dtype_builtin2", ret_dtype=FP64) + + # lazy=True must not defer the validation: a bad combination fails at the + # registration site, not at first attribute touch of the delayed op. + with pytest.raises(ValueError, match="parameterized=True"): + UnaryOp.register_new( + "_ret_dtype_lazy_param", + lambda x: x, + parameterized=True, + is_udt=True, + lazy=True, + ret_dtype=FP64, + ) + + # A parameterized operator builds its function when called; ret_dtype + # belongs to the register call for that function. + with pytest.raises(ValueError, match="does not work with parameterized=True"): + UnaryOp.register_anonymous( + lambda t=1: (lambda x: x + t), + "_ret_dtype_param", + parameterized=True, + is_udt=True, + ret_dtype=FP64, + ) + + # SelectOp takes no ret_dtype: GraphBLAS fixes its return type to BOOL. + with pytest.raises(TypeError, match="ret_dtype"): + SelectOp.register_anonymous( + lambda x, i, j, t: x > t, "_ret_dtype_select", is_udt=True, ret_dtype=BOOL + ) + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_ret_dtype_index_ops(): + """``ret_dtype`` reaches the IndexUnaryOp and IndexBinaryOp compile paths too.""" + # Shape (11,) is used nowhere else: anonymous UDTs share one DataType per + # np.dtype, so a shape reused across tests would inherit whichever JIT C + # state was frozen first under random test ordering. + in11 = dtypes.register_anonymous(np.dtype((np.float64, (11,))), "_RetDIdx11") + out2 = dtypes.register_anonymous(np.dtype((np.float64, (2,))), "_RetDIdx2") + + def _head_plus_row(x, i, j, t): # pragma: no cover (numba) + return x[:2] + i + + iu = IndexUnaryOp.register_anonymous( + _head_plus_row, "_ret_dtype_indexunary", is_udt=True, ret_dtype=out2 + ) + assert iu[in11, INT64].return_type is out2 + + if lib.__dict__.get("GxB_IndexBinaryOp_new") is not None: + from graphblas.core.operator import IndexBinaryOp + + def _head_sum(x, ix, jx, y, iy, jy, theta): # pragma: no cover (numba) + return x[:2] + y[:2] + + ib = IndexBinaryOp.register_anonymous( + _head_sum, "_ret_dtype_indexbinary", is_udt=True, ret_dtype=out2 + ) + assert ib[in11, in11].return_type is out2 + + +@pytest.mark.skipif("not supports_udfs") +@pytest.mark.slow +def test_udt_ret_dtype_still_shape_checked(): + """The registration-time shape probe checks against the declared ``ret_dtype``. + + Short-circuiting the return-type inference must not take the probe with + it. Declaring the output type makes the check worth more, not less: the + inferred case can only ever compare a UDF against a type derived from + what it returned, while a declared type is an independent claim the UDF + can contradict. + """ + # Shape (14,) is used nowhere else: anonymous UDTs share one DataType per + # np.dtype, so a shape reused across tests would inherit whichever JIT C + # state was frozen first under random test ordering. + in14 = dtypes.register_anonymous(np.dtype((np.float64, (14,))), "_RetDPrb14") + out2 = dtypes.register_anonymous(np.dtype((np.float64, (2,))), "_RetDPrb2") + + def _three(x): # pragma: no cover (numba) + return x[:3] + + # Three elements cannot fill a declared two-element output. + op_ = UnaryOp.register_anonymous(_three, "_ret_dtype_prb_bad", is_udt=True, ret_dtype=out2) + with pytest.raises(UdfParseError, match=r"shape \(3,\).*_RetDPrb2 elements are \(2,\)"): + op_[in14] + + # The check is fit-by-broadcast, not equality: a one-element return + # legitimately fills every slot of the declared element. + def _one(x): # pragma: no cover (numba) + return x[:1] + + ok = UnaryOp.register_anonymous(_one, "_ret_dtype_prb_bcast", is_udt=True, ret_dtype=out2) + assert ok[in14].return_type is out2 + + # The record half of the probe checks the declared type's array leaves. + rec = dtypes.register_anonymous( + np.dtype([("v", np.float64, (4,)), ("n", np.int64)], align=True), "_RetDPrbRec" + ) + + def _short_leaf(x, y): # pragma: no cover (numba) + return (x[:2], 1) + + recop = BinaryOp.register_anonymous( + _short_leaf, "_ret_dtype_prb_rec", is_udt=True, ret_dtype=rec + ) + with pytest.raises(UdfParseError, match=r"shape \(2,\) for field \['v'\] of _RetDPrbRec"): + recop[in14, in14] + + +def _ret_dtype_pickle_udf(x): # pragma: no cover (numba) + return x[:5] + + +_RET_DTYPE_PICKLE_WRITER = """ +import pickle +import sys + +import numpy as np + +import graphblas as gb +from graphblas.core.operator.unary import UnaryOp +from graphblas.tests.test_op import _ret_dtype_pickle_udf + +print("package: " + gb.__file__) +five = gb.dtypes.register_anonymous(np.dtype((np.float32, (5,))), "_RetPickle5") +UnaryOp.register_new("_ret_dtype_pickled", _ret_dtype_pickle_udf, is_udt=True, ret_dtype=five) +anon = UnaryOp.register_anonymous( + _ret_dtype_pickle_udf, "_ret_dtype_pickled_anon", is_udt=True, ret_dtype=five +) +with open(sys.argv[1], "wb") as f: + pickle.dump((gb.unary._ret_dtype_pickled, anon), f) +print("wrote: ok") +""" + + +_RET_DTYPE_PICKLE_READER = """ +import pickle +import sys + +import numpy as np + +import graphblas as gb + +print("package: " + gb.__file__) +with open(sys.argv[1], "rb") as f: + named, anon = pickle.load(f) +ten = gb.dtypes.register_anonymous(np.dtype((np.float32, (10,))), "_RetPickle10") +print("named shape: " + str(named[ten].return_type.np_type.subdtype[1])) +print("anon shape: " + str(anon[ten].return_type.np_type.subdtype[1])) +""" + + +@pytest.mark.skipif("not supports_udfs") +def test_udt_ret_dtype_survives_pickle(tmp_path): + # In-process unpickling takes the _find shortcut and returns the very + # same object, so only a cross-process round trip exercises what pickle + # exists for: the reduce tuple must carry ret_dtype, or the reconstructed + # op re-registers without it and fails its own shape probe. Both sides + # run in subprocesses so the named registration never lands in this + # process's operator namespace (test_operator_types enumerates it). + payload = tmp_path / "ops.pkl" + repo_root = Path(__file__).resolve().parents[2] + env = dict(os.environ, PYTHONPATH=str(repo_root)) + for probe, checks in ( + (_RET_DTYPE_PICKLE_WRITER, {"wrote": "ok"}), + (_RET_DTYPE_PICKLE_READER, {"named shape": "(5,)", "anon shape": "(5,)"}), + ): + result = subprocess.run( + [sys.executable, "-c", probe, str(payload)], + capture_output=True, + text=True, + check=False, + env=env, + ) + report = f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + assert result.returncode == 0, report + lines = dict(line.split(": ", 1) for line in result.stdout.splitlines() if ": " in line) + assert lines.get("package") == str(repo_root / "graphblas" / "__init__.py"), report + for key, want in checks.items(): + assert lines.get(key) == want, report diff --git a/graphblas/tests/test_scalar.py b/graphblas/tests/test_scalar.py index 49e7221b4..ac45a37f1 100644 --- a/graphblas/tests/test_scalar.py +++ b/graphblas/tests/test_scalar.py @@ -429,6 +429,7 @@ def test_expr_is_like_scalar(s): expected = { "__call__", "__del__", + "__getattr__", "__imatmul__", "__lshift__", "_carg", @@ -473,6 +474,7 @@ def test_index_expr_is_like_scalar(s): # Should we make any of these raise informative errors? expected = { "__del__", + "__getattr__", "__imatmul__", "_carg", "_deserialize", @@ -670,6 +672,59 @@ def test_get(s): assert s.get("mittens") == "mittens" +@autocompute +def test_index_expr_value_fast_path(): + """v[i].value / float(v[i]) resolve with a single extract (see automethods). + + The read must match get() / .new().value exactly and the UDT + Recorder + fallbacks must keep working; autocompute=False still raises (tested below). + """ + v = Vector.from_coo([0, 2, 5], [1.5, 2.5, 3.5], size=10) + assert v[2].value == 2.5 + assert type(v[2].value) is float + assert float(v[2]) == 2.5 + assert int(v[2]) == 2 # Scalar.__int__ truncates + assert v[2].is_empty is False + assert v[4].value is None # miss + assert v[4].is_empty is True + with pytest.raises(TypeError): + float(v[4]) # float(None) + assert v[2].value == v[2].new().value == v.get(2) + + A = Matrix.from_coo([0, 1, 4], [1, 2, 6], [1.5, 2.5, 3.5], nrows=5, ncols=7) + assert A[1, 2].value == 2.5 + assert A.T[2, 1].value == 2.5 + assert A[0, 0].value is None + + # UDT reads defer to `.new()` (numpy conversion in Scalar.value) + udt = dtypes.register_anonymous( + np.dtype([("sx", np.int64), ("sy", np.float64)]), "_ScalarValueFastUdt" + ) + u = Vector(udt, 3) + u[1] = (7, 2.5) + assert u[1].value["sx"] == 7 + assert u[0].value is None + + # A Recorder must still observe the extract (fall-back path) + from graphblas.core.recorder import Recorder + + v2 = Vector.from_coo([0], [1.5], size=4) + with Recorder() as rec: + assert v2[0].value == 1.5 + assert "extractElement" in "".join(rec.data) + + +def test_index_expr_value_autocompute_false(): + """The value path stays gated on autocompute; .new() remains the escape hatch.""" + v = Vector.from_coo([0, 2, 5], [1.5, 2.5, 3.5], size=10) + with gb.config.set(autocompute=False): + with pytest.raises(TypeError): + v[2].value + with pytest.raises(TypeError): + float(v[2]) + assert v[2].new().value == 2.5 + + def test_ss_descriptors(s): v = Vector.from_coo([0, 2], [10, 20]) if suitesparse: diff --git a/graphblas/tests/test_ss_refcount.py b/graphblas/tests/test_ss_refcount.py new file mode 100644 index 000000000..d37a30014 --- /dev/null +++ b/graphblas/tests/test_ss_refcount.py @@ -0,0 +1,144 @@ +"""Regression tests for gh-559: Matrix/Vector must not be born in a reference cycle. + +Historically each Matrix and Vector stored ``self.ss = ss(self)``, and the ss +object (and its config) held ``_parent`` back to the object, so instances died +only via the cyclic garbage collector, not by reference counting. For large +matrices in a tight loop this deferred the release of the C-side GrB buffer +until gc happened to run. ``.ss`` is now built lazily per access and stored +nowhere, so the cycle never forms. These tests pin that behavior; they are +suitesparse-only because ``.ss`` exists only on that backend. +""" + +import gc +import weakref + +import numpy as np +import pytest + +from graphblas import Matrix, Vector, backend, semiring +from graphblas.core.ss.matrix import ss as matrix_ss_class +from graphblas.core.ss.vector import ss as vector_ss_class + +if backend != "suitesparse": + pytest.skip("A.ss only available with suitesparse backend", allow_module_level=True) + + +def _dies_by_reference_count(factory): + """True if the object dies the instant its last strong ref drops, with gc off.""" + gc.collect() + was_enabled = gc.isenabled() + gc.disable() + try: + obj = factory() + wref = weakref.ref(obj) + del obj + return wref() is None + finally: + if was_enabled: + gc.enable() + + +def test_matrix_dies_by_reference_count(): + assert _dies_by_reference_count(lambda: Matrix(float, 5, 5)) + + +def test_vector_dies_by_reference_count(): + assert _dies_by_reference_count(lambda: Vector(float, 5)) + + +def test_mxm_result_dies_by_reference_count(): + A = Matrix.from_dense(np.arange(9.0).reshape(3, 3) + 1) + B = Matrix.from_dense(np.arange(9.0).reshape(3, 3) + 1) + assert _dies_by_reference_count(lambda: A.mxm(B, semiring.min_plus).new()) + + +def test_ss_namespace_is_functional(): + A = Matrix.from_coo([0, 1, 2], [0, 1, 2], [1.0, 2.0, 3.0], nrows=3, ncols=3) + # Built fresh each access (nothing is stored on the instance). + assert A.ss is not A.ss + # Introspection still works through a fresh access. + assert A.ss.nbytes > 0 + matrix_formats = { + "csr", + "csc", + "hypercsr", + "hypercsc", + "bitmapr", + "bitmapc", + "fullr", + "fullc", + "coor", + "cooc", + } + assert A.ss.export()["format"] in matrix_formats + # Config get then set then get, each through an independent `.ss`. + assert A.ss.config["format"] in {"by_row", "by_col"} + A.ss.config["format"] = "by_col" + assert A.ss.config["format"] == "by_col" + + v = Vector.from_coo([0, 2], [1.0, 3.0], size=4) + assert v.ss is not v.ss + assert v.ss.nbytes > 0 + assert v.ss.export()["format"] in {"sparse", "bitmap", "full"} + + +def test_ss_class_access_returns_namespace_class(): + # Class-level access must still yield the ss class so its import_* classmethods work. + assert Matrix.ss is matrix_ss_class + assert Vector.ss is vector_ss_class + assert hasattr(Matrix.ss, "import_any") + assert hasattr(Vector.ss, "import_any") + + +def test_ss_attribute_is_read_only(): + A = Matrix(float, 3, 3) + with pytest.raises(AttributeError): + A.ss = 5 + v = Vector(float, 3) + with pytest.raises(AttributeError): + v.ss = 5 + + +def test_views_have_working_ss(): + # A single-column Matrix cast to a Vector (_as_vector) is a view with _parent set. + A = Matrix.from_coo([0, 1], [0, 0], [1.0, 2.0], nrows=3, ncols=1) + v = A._as_vector() + assert v._parent is A + assert v.ss.nbytes > 0 + # A Vector cast to a Matrix (_as_matrix) is likewise a view. + w = Vector.from_coo([0, 2], [1.0, 3.0], size=4) + M = w._as_matrix() + assert M._parent is w + assert M.ss.nbytes > 0 + + +def test_batched_mxm_loop_does_not_accumulate_matrices(): + # gh-559: with the cyclic collector switched off, a batched + # mxm -> to_dense -> discard loop must not pile up Matrix objects. + rng = np.random.default_rng(0) + A = Matrix.from_dense(rng.random((16, 8)) + 0.1) + B = Matrix.from_dense(rng.random((8, 24)) + 0.1) + + def live_matrices(): + return sum(1 for o in gc.get_objects() if type(o) is Matrix) + + gc.collect() + was_enabled = gc.isenabled() + gc.disable() + try: + # Warm up one iteration so any one-time caches are populated first. + C = Matrix(float, 16, 24) + C << A.mxm(B, semiring.min_plus) + C.to_dense(0.0) + del C + baseline = live_matrices() + for _ in range(50): + C = Matrix(float, 16, 24) + C << A.mxm(B, semiring.min_plus) + C.to_dense(0.0) + del C + # Without the fix this would be baseline + 50. + assert live_matrices() <= baseline + finally: + if was_enabled: + gc.enable() diff --git a/graphblas/tests/test_ss_utils.py b/graphblas/tests/test_ss_utils.py index 40774186c..30a664742 100644 --- a/graphblas/tests/test_ss_utils.py +++ b/graphblas/tests/test_ss_utils.py @@ -237,6 +237,80 @@ def test_global_config(): assert "format" in repr(config) +def test_global_config_key_completions(): + # IPython offers these for `config[]`, so they must be the keys that + # actually resolve, not the attributes of the mapping object. + config = gb.ss.config + completions = config._ipython_key_completions_() + assert set(completions) == set(config._options) + for key in completions: + config[key] + # About aliases the same hook onto its own __iter__ (gb.ss.about[]) + about = gb.ss.about + completions = about._ipython_key_completions_() + assert set(completions) == set(about) + for key in completions: + about[key] + + +def test_global_config_attribute_assignment_raises(): + # Attribute assignment used to create a dead instance attribute and leave + # the real config unchanged; now it raises and points at item assignment. + config = gb.ss.config + before = config["nthreads"] + with pytest.raises(AttributeError, match="item assignment"): + config.nthreads = before + 1 + assert config["nthreads"] == before + assert "nthreads" not in config.__dict__ + with pytest.raises(AttributeError, match="Unknown config option 'nthread'"): + config.nthread = before + 1 + # Item assignment is the supported write and still works + config["nthreads"] = before + 1 + assert config["nthreads"] == before + 1 + config["nthreads"] = before + assert config["nthreads"] == before + # Reads, iteration, and key completions are unaffected by the guard + assert dict(config) == {k: config[k] for k in config} + assert set(config._ipython_key_completions_()) == set(config._options) + + +def test_object_config_attribute_assignment_raises(): + v = Vector(int, 3) + with pytest.raises(AttributeError, match="item assignment"): + v.ss.config.sparsity_control = "bitmap" + assert v.ss.config["sparsity_control"] == {"auto"} + with pytest.raises(AttributeError, match="read-only"): + v.ss.config.sparsity_status = "bitmap" + A = Matrix(int, 2, 2) + with pytest.raises(AttributeError, match="item assignment"): + A.ss.config.format = "by_col" + assert A.ss.config["format"] == "by_row" + # Item assignment is the supported write and still works + v.ss.config["sparsity_control"] = "bitmap" + assert v.ss.config["sparsity_control"] == {"bitmap"} + v.ss.config["sparsity_control"] = "auto" + assert v.ss.config["sparsity_control"] == {"auto"} + + +def test_about_attribute_assignment_raises(): + about = gb.ss.about + mode = about["mode"] + with pytest.raises(AttributeError, match="read-only"): + about.mode = "junk" + assert about["mode"] == mode + assert "mode" not in about.__dict__ + with pytest.raises(AttributeError, match="Unknown About option"): + about.junkattr = 1 + assert dict(about) == {k: about[k] for k in about} + + +@pytest.mark.skipif("gb.core.ss._IS_SSGB7") +def test_context_attribute_assignment_raises(): + context = gb.ss.Context(engage=False) + with pytest.raises(AttributeError, match="item assignment"): + context.nthreads = 4 + + @pytest.mark.skipif("gb.core.ss._IS_SSGB7") def test_context(): context = gb.ss.Context() diff --git a/graphblas/tests/test_ssjit.py b/graphblas/tests/test_ssjit.py index fc174c768..b75c4ac36 100644 --- a/graphblas/tests/test_ssjit.py +++ b/graphblas/tests/test_ssjit.py @@ -735,41 +735,63 @@ def test_floordiv_udt_jit_matches_python_semantics(): @pytest.mark.skipif("not supports_udfs") -def test_min_max_udt_jit_propagates_nan(): - """``binary.min``/``max`` on a float UDT must propagate NaN like Python/numba. - - Regression: the JIT codegen used to emit ``(a < b ? a : b)``, which - silently swallows NaN to the right-hand side. Python ``min(a, b)`` (and - numba's ``min``) returns ``a`` when neither comparison is true (NaN - involved), so ``min(NaN, 1.0) == NaN`` and ``min(1.0, NaN) == 1.0``. - The fix swaps the ternary to ``(b < a ? b : a)``. +def test_min_max_udt_jit_calls_fmin_and_ignores_nan(): + """The JIT kernel for ``binary.min`` on a float UDT must call C ``fmin``. + + ``GrB_MIN_FP64`` is C99 ``fmin``, which ignores a NaN operand from either + side. ``binary.min`` has to mean the same thing when it is typed for a + UDT as when it is typed for FP64, so the kernel calls ``fmin`` rather + than deciding NaN with a comparison. Two earlier spellings decided it, + in opposite directions: ``(a < b ? a : b)`` dropped a NaN on the right, + and ``(b < a ? b : a)`` dropped one on the left to agree with Python's + builtin ``min``, which the cfunc path was reaching by accident. Either + way the answer turned on which operand the NaN arrived on. + + Integer fields keep the comparison: they have no NaN to order, and it + saves a conversion through ``double`` per element. """ - if _IS_SSGB7: - pytest.skip("JIT requires SuiteSparse:GraphBLAS >= 8") + if not _has_jit_set: + pytest.skip("jit_c_source introspection requires SuiteSparse:GraphBLAS >= 9") _require_jit_on() # Field names unique to this test; see floordiv test for the cache rationale. udt = dtypes.register_anonymous( - np.dtype([("nan_a", np.float64), ("nan_b", np.float64)]), "_NanJitMM" + np.dtype([("nan_a", np.float64), ("nan_b", np.float32), ("nan_c", np.int32)]), "_NanJitMM" ) + csrc = binary.min[udt].jit_c_source + assert "fmin((x->nan_a), (y->nan_a))" in csrc, csrc + assert "fminf((x->nan_b), (y->nan_b))" in csrc, csrc + assert "((x->nan_c) < (y->nan_c) ? (x->nan_c) : (y->nan_c))" in csrc, csrc + assert "fmax((x->nan_a), (y->nan_a))" in binary.max[udt].jit_c_source + N = 100 v = gb.Vector(udt, N) u = gb.Vector(udt, N) nan = float("nan") for i in range(N): # field nan_a: NaN on the left; field nan_b: NaN on the right at odd indices. - v[i] = (nan, 2.0 + i) - u[i] = (1.0 + i, nan if i % 2 else 3.0 + i) + v[i] = (nan, 2.0 + i, i) + u[i] = (1.0 + i, nan if i % 2 else 3.0 + i, 2 * i) w = v.ewise_mult(u, binary.min).new() - assert np.isnan(w[0].new().value[0]) # min(NaN, 1.0) -> NaN - assert w[1].new().value[1] == 3.0 # min(3.0, NaN) -> 3.0 (NaN swallowed) + assert w[0].new().value[0] == 1.0 # min(NaN, 1.0) -> 1.0 + assert w[1].new().value[1] == 3.0 # min(3.0, NaN) -> 3.0 assert w[2].new().value[1] == 4.0 # min(4.0, 5.0) -> 4.0 (normal case) + assert w[3].new().value[2] == 3 # integer field is unaffected w = v.ewise_mult(u, binary.max).new() - assert np.isnan(w[0].new().value[0]) # max(NaN, 1.0) -> NaN - assert w[1].new().value[1] == 3.0 # max(3.0, NaN) -> 3.0 (NaN swallowed) + assert w[0].new().value[0] == 1.0 # max(NaN, 1.0) -> 1.0 + assert w[1].new().value[1] == 3.0 # max(3.0, NaN) -> 3.0 assert w[2].new().value[1] == 5.0 # max(4.0, 5.0) -> 5.0 (normal case) + assert w[3].new().value[2] == 6 # integer field is unaffected + + # Both operands NaN is the one case where a NaN survives, for min and max + # alike, and it is the only case ``fmin`` has no non-NaN answer for. + nan_only = gb.Vector(udt, N) + for i in range(N): + nan_only[i] = (nan, np.float32(i), i) + w = nan_only.ewise_mult(nan_only, binary.min).new() + assert np.isnan(w[0].new().value[0]) @pytest.mark.skipif("not supports_udfs") diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py index 52c382742..d95ef386e 100644 --- a/graphblas/tests/test_vector.py +++ b/graphblas/tests/test_vector.py @@ -728,6 +728,30 @@ def test_apply_indexunary(v): v.apply(indexunary.valueeq, left=s2) +def test_indexunary_helpers(v): + # indexunary.value/row/column mirror select's helpers (GH #239), returning + # apply expressions (BOOL) rather than select expressions. + assert indexunary.value(v == 1).new().isequal(v.apply(indexunary.valueeq, 1).new()) + assert indexunary.value(v >= 2).new().isequal(v.apply(indexunary.valuege, 2).new()) + # Index comparisons on a Vector go through the row helper (indexle is rowle). + assert indexunary.row(v < 4).new().isequal(v.apply(indexunary.rowle, 3).new()) + assert indexunary.row(v >= 4).new().isequal(v.apply(indexunary.rowgt, 3).new()) + assert indexunary.value(v == 1).new().dtype == dtypes.BOOL + assert indexunary.row(v < 4).new().dtype == dtypes.BOOL + # Scalar expression path resolves to a ScalarExpression, like select.value. + s = Scalar.from_value(1) + assert indexunary.value(s < 10).new() == s.apply(indexunary.valuelt, 10).new() + # `index` is intentionally not a helper: it remains the rowindex op alias, so + # `v.apply(indexunary.index)` keeps working. + assert indexunary.index is indexunary.rowindex + # A value mask reads naturally through the value helper. + assert ( + indexunary.value(v != False).new().isequal(indexunary.valuene(v, False).new()) # noqa: E712 + ) + with pytest.raises(TypeError, match="indexunary.value"): + indexunary.value(v) + + def test_select(v): result = Vector.from_coo([1, 3], [1, 1], size=7) w1 = v.select(select.valueeq, 1).new() @@ -1658,6 +1682,7 @@ def test_expr_is_like_vector(v): "__call__", "__del__", "__delitem__", + "__getattr__", "__lshift__", "__setitem__", "_assign_element", @@ -1708,6 +1733,7 @@ def test_index_expr_is_like_vector(v): expected = { "__del__", "__delitem__", + "__getattr__", "__setitem__", "_assign_element", "_delete_element", @@ -2688,3 +2714,23 @@ def test_subarray_dtypes(): assert full1.isequal(full2, check_dtype=True) full2 = Vector.ss.import_bitmap(values=a, bitmap=[True, True, True]) assert full1.isequal(full2, check_dtype=True) + + +def test_constructor_rejects_arraylike_first_arg(): + # The first positional arg is the dtype; passing data instead used to raise a + # confusing "Unknown dtype" ValueError. It should now raise TypeError pointing + # at the from_* constructors. + with pytest.raises(TypeError, match="Vector.*dtype.*from_coo.*from_dense"): + Vector([1, 2, 3]) + with pytest.raises(TypeError, match="Vector.*expects a dtype"): + Vector((1, 2, 3)) + with pytest.raises(TypeError, match="Vector.*expects a dtype"): + Vector(np.array([1, 2, 3])) + # Valid dtype-first signatures still work, including list/tuple dtype specs + assert Vector(int, size=3).dtype == dtypes.INT64 + assert Vector("INT64", size=3).dtype == dtypes.INT64 + assert Vector(np.dtype("int64"), size=3).dtype == dtypes.INT64 + assert Vector([("x", "i8"), ("y", "f8")], size=3).dtype._is_udt + # A non-array-like bad dtype keeps the original ValueError + with pytest.raises(ValueError, match="Unknown dtype"): + Vector("not_a_dtype", size=3) diff --git a/graphblas/tests/test_viz.py b/graphblas/tests/test_viz.py new file mode 100644 index 000000000..232d01ce3 --- /dev/null +++ b/graphblas/tests/test_viz.py @@ -0,0 +1,258 @@ +"""Smoke tests for graphblas.viz. + +The viz module is optional-dependency heavy (matplotlib, networkx, scipy for +``spy``/``draw``; datashader + holoviews + hvplot + bokeh + pandas for +``datashade``). These tests only check that each public function runs end to end +under the headless Agg backend and populates a figure/returns an object; they do +not assert on pixel output. Anything missing is skipped, not failed, so a +minimal-dependency CI run sees clean skips. +""" + +import importlib +import inspect +import math +import warnings + +import pytest + +from graphblas import Matrix, Vector, viz + + +def _importorskip(modname): + """Skip when an optional dependency is missing *or* installed but unusable. + + From pytest 9.1 on, ``pytest.importorskip`` counts only ``ModuleNotFoundError`` + as "missing". A dependency that is installed but cannot run raises a plain + ``ImportError`` instead (matplotlib does exactly that when numpy is older than + it supports), which escapes ``importorskip`` and aborts collection for the + whole session. pytest's ``exc_type`` argument covers that, but it only exists + in pytest >=8.2 and this project supports pytest >=6.2, so do the import here + and hand the already-imported module to pytest. + """ + try: + with warnings.catch_warnings(): + # ``importorskip`` ignores warnings while importing; match it, or + # ``filterwarnings = error`` would fail on whatever an optional + # dependency happens to emit at import time. + warnings.simplefilter("ignore") + importlib.import_module(modname) + except ImportError as exc: + pytest.skip(f"could not import {modname!r}: {exc}", allow_module_level=True) + return pytest.importorskip(modname) + + +# Skip the whole module if matplotlib is absent (draw and spy both need it). +# Set the backend to Agg before pyplot is imported so no display is required. +mpl = _importorskip("matplotlib") +mpl.use("Agg") +plt = _importorskip("matplotlib.pyplot") + + +@pytest.fixture(autouse=True) +def _close_figures(): + # Close every figure after each test to avoid matplotlib's + # "More than 20 figures have been opened" warning (which the project's + # ``filterwarnings = error`` config would turn into a failure). + yield + plt.close("all") + + +def square_matrix(): + # Small square adjacency matrix with distinct weights. + return Matrix.from_coo([0, 0, 1, 2], [1, 2, 2, 0], [1.0, 2.0, 3.0, 4.0], nrows=3, ncols=3) + + +def test_spy_default(): + _importorskip("scipy.sparse") + A = square_matrix() + fig = viz.spy(A, show=False) + assert isinstance(fig, mpl.figure.Figure) + assert fig.axes, "spy should populate at least one Axes" + # matplotlib's Axes.spy draws the pattern as a single markered Line2D. + assert fig.axes[0].lines, "spy should plot the sparsity markers" + + +def test_spy_centered(): + # centered=True skips the tick-offset fixup branch. + _importorskip("scipy.sparse") + A = square_matrix() + fig = viz.spy(A, show=False, centered=True) + assert isinstance(fig, mpl.figure.Figure) + assert fig.axes[0].lines + + +def test_spy_with_axes(): + # Passing an explicit Axes exercises the ``axes is not None`` branch, + # including the auto-markersize path (which once raised NameError here). + _importorskip("scipy.sparse") + A = square_matrix() + fig = mpl.figure.Figure() + axes = fig.subplots() + result = viz.spy(A, show=False, axes=axes) + assert result is fig + assert axes.lines + + +def test_spy_with_figure(): + # Passing an explicit Figure (no Axes) once raised NameError; spy should + # create the Axes on the given figure and return that same figure. + _importorskip("scipy.sparse") + A = square_matrix() + fig = mpl.figure.Figure() + result = viz.spy(A, show=False, figure=fig) + assert result is fig + assert fig.axes + assert fig.axes[0].lines + + +@pytest.mark.filterwarnings("ignore:FigureCanvasAgg is non-interactive") +def test_draw(): + # draw() renders onto the current pyplot Axes via networkx and calls + # plt.show(); on Agg that show() emits the non-interactive UserWarning, + # which we ignore here. + _importorskip("networkx") + _importorskip("scipy.sparse") + A = square_matrix() + viz.draw(A) + axes = plt.gcf().get_axes() + assert axes, "draw should populate the current figure" + ax = axes[0] + # Nodes render as patches/collections and labels as texts. + assert ax.collections or ax.patches + assert ax.texts, "draw should render node/edge labels" + + +def test_draw_rejects_non_matrix(): + _importorskip("networkx") + v = Vector.from_coo([0, 1, 2], [1.0, 2.0, 3.0]) + with pytest.raises(TypeError, match="Can only draw a Matrix"): + viz.draw(v) + + +@pytest.mark.filterwarnings("ignore:FigureCanvasAgg is non-interactive") +def test_draw_reciprocal_edges_both_labels_visible(): + # Regression for gh-474: reciprocal directed edges (0->1 and 1->0) used to be + # drawn as coincident straight lines, so one weight hid the other. draw() now + # curves reciprocal pairs; both weights must appear at distinct positions. + nx = _importorskip("networkx") + _importorskip("scipy.sparse") + # draw() only curves reciprocal pairs when networkx can place edge labels along + # the curve; without that it deliberately draws every edge straight, and the two + # labels then coincide (which test_draw_without_networkx_curved_label_support + # covers). Ask the same question draw() asks, so the two cannot drift apart. + if "connectionstyle" not in inspect.signature(nx.draw_networkx_edge_labels).parameters: + pytest.skip("networkx <3.3: draw_networkx_edge_labels has no connectionstyle") + M = Matrix.from_coo([0, 1], [1, 0], [10, 20], nrows=2, ncols=2) + viz.draw(M) + ax = plt.gcf().get_axes()[0] + + weight_labels = [t for t in ax.texts if t.get_text() in {"10", "20"}] + assert {t.get_text() for t in weight_labels} == {"10", "20"}, "both weights must be drawn" + assert len(weight_labels) == 2 + + # Old behavior placed both labels on the shared straight-line midpoint. + # networkx returns two anchors there that differ only by floating-point + # noise (order 1e-6), so an exact ``!=`` comparison passes even when the + # labels sit on top of each other. Require a separation that is a real + # fraction of the distance between the two nodes instead. + node_positions = [t.get_position() for t in ax.texts if t.get_text() in {"0", "1"}] + assert len(node_positions) == 2 + edge_length = math.dist(*node_positions) + separation = math.dist(*(t.get_position() for t in weight_labels)) + assert separation > 0.001 * edge_length, "reciprocal edge labels still overlap" + + +@pytest.mark.filterwarnings("ignore:FigureCanvasAgg is non-interactive") +def test_draw_without_networkx_curved_label_support(monkeypatch): + # draw_networkx_edge_labels gained connectionstyle in networkx 3.3, and the + # project supports >=2.8. Standing in a pre-3.3 signature must not raise; the + # gh-474 curving is skipped and every edge renders straight, as it did before. + nx = _importorskip("networkx") + _importorskip("scipy.sparse") + real = nx.draw_networkx_edge_labels + + def pre_33_draw_networkx_edge_labels(g, pos, edge_labels=None, **kwargs): + if "connectionstyle" in kwargs: + raise TypeError( + "draw_networkx_edge_labels() got an unexpected keyword argument " + "'connectionstyle'" + ) + return real(g, pos, edge_labels=edge_labels, **kwargs) + + monkeypatch.setattr(nx, "draw_networkx_edge_labels", pre_33_draw_networkx_edge_labels) + M = Matrix.from_coo([0, 1], [1, 0], [10, 20], nrows=2, ncols=2) + viz.draw(M) + ax = plt.gcf().get_axes()[0] + assert {t.get_text() for t in ax.texts if t.get_text() in {"10", "20"}} == {"10", "20"} + + +def _import_datashade_deps(): + for name in ("numpy", "pandas", "datashader", "holoviews", "hvplot", "bokeh"): + _importorskip(name) + + +def test_datashade_single(): + _import_datashade_deps() + import holoviews as hv + + A = square_matrix() + obj = viz.datashade(A) + assert obj is not None + assert isinstance(obj, hv.core.dimension.Dimensioned) + + +def test_datashade_agg_list(): + # A flat list of aggregators produces one row of linked plots. + _import_datashade_deps() + import holoviews as hv + + A = square_matrix() + layout = viz.datashade(A, agg=["count", "sum"]) + assert isinstance(layout, hv.Layout) + + +def test_datashade_agg_grid(): + # A list-of-lists produces a 2d grid of linked plots. + _import_datashade_deps() + import holoviews as hv + + A = square_matrix() + layout = viz.datashade(A, agg=[["count", "sum"], ["min", "max"]]) + assert isinstance(layout, hv.Layout) + + +def test_datashade_empty_agg(): + # An empty aggregator list is a no-op that returns None. + _import_datashade_deps() + A = square_matrix() + assert viz.datashade(A, agg=[]) is None + + +def test_datashade_positions_match_spy(): + # Regression for gh-473: element (row=r, col=c) must render centered on the + # integer tick pair (col, row), the same convention ``spy`` uses. We check + # the datashader aggregation directly (no display) over the limits the + # interactive path uses, at one pixel per matrix cell. + _import_datashade_deps() + import datashader as ds + import numpy as np + + # Non-square (3x4) with distinct row/col so a row<->col swap would show. + M = Matrix.from_coo([0, 0, 2], [1, 3, 3], [1.0, 1.0, 1.0], nrows=3, ncols=4) + df = viz._matrix_to_dataframe(M) + xlim, ylim = viz._cell_centered_limits(M) + assert xlim == (-0.5, M.ncols - 0.5) + assert ylim == (-0.5, M.nrows - 0.5) + + canvas = ds.Canvas(plot_width=M.ncols, plot_height=M.nrows, x_range=xlim, y_range=ylim) + agg = canvas.points(df, "col", "row", ds.count()) + + # Pixel centers land on integers, so ticks label the cells they sit on. + assert agg.coords["col"].values.tolist() == [0.0, 1.0, 2.0, 3.0] + assert agg.coords["row"].values.tolist() == [0.0, 1.0, 2.0] + + # Counts are nonzero exactly at the (row, col) indices of the elements. + xs = agg.coords["col"].values + ys = agg.coords["row"].values + nonzero = {(round(float(ys[i])), round(float(xs[j]))) for i, j in np.argwhere(agg.values > 0)} + assert nonzero == {(0, 1), (0, 3), (2, 3)} diff --git a/graphblas/unary/__init__.py b/graphblas/unary/__init__.py index b83ea3b8b..e48a6a387 100644 --- a/graphblas/unary/__init__.py +++ b/graphblas/unary/__init__.py @@ -42,7 +42,9 @@ def __getattr__(key): ss = import_module(".ss", __name__) globals()["ss"] = ss return ss - raise AttributeError(f"module {__name__!r} has no attribute {key!r}") + from ..core.utils import _module_attr_error + + raise _module_attr_error(__name__, key, __dir__()) from ..core import operator # noqa: E402 isort:skip diff --git a/graphblas/viz.py b/graphblas/viz.py index 8e2a53228..8e685b374 100644 --- a/graphblas/viz.py +++ b/graphblas/viz.py @@ -42,12 +42,17 @@ def _get_imports(names, within): return rv -def draw(m): # pragma: no cover +def draw(m): """Draw a square adjacency Matrix as a graph. Requires `networkx `_ and `matplotlib `_ to be installed. + Reciprocal directed edges (``u -> v`` and ``v -> u``) are drawn as curves so + both arrows and both edge weights stay visible; all other edges are straight. + Curving them needs networkx 3.3 or newer; with older versions every edge is + drawn straight. + Example output: .. image:: /_static/img/draw-example.png @@ -59,9 +64,44 @@ def draw(m): # pragma: no cover g = to_networkx(m) pos = nx.spring_layout(g) - edge_labels = {(i, j): d["weight"] for i, j, d in g.edges(data=True)} - nx.draw_networkx(g, pos, node_color="red", node_size=500) - nx.draw_networkx_edge_labels(g, pos, edge_labels=edge_labels) + node_size = 500 + nx.draw_networkx_nodes(g, pos, node_color="red", node_size=node_size) + nx.draw_networkx_labels(g, pos) + + # A reciprocal pair (u -> v and v -> u) drawn as two straight lines coincides, + # hiding one edge's weight (python-graphblas #474). Curving both edges makes + # each bend toward its own side, so both arrows and both labels stay visible + # and attributable. Self-loops (u == v) are not reciprocal; leave them straight. + # + # networkx only learned to place edge labels along a curve in 3.3, and we + # support >=2.8, so fall back to the previous straight rendering without it. + # Curving the edges but not the labels would be worse than not curving at all: + # the labels would sit back on the shared chord midpoint, which is the overlap + # #474 is about, and they would no longer track their arrows. Check the + # parameter rather than pin a version. + import inspect + + if "connectionstyle" in inspect.signature(nx.draw_networkx_edge_labels).parameters: + curved = {(u, v) for u, v in g.edges if u != v and g.has_edge(v, u)} + else: + curved = set() + straight = [e for e in g.edges if e not in curved] + connectionstyle = "arc3,rad=0.1" + + def _edge_labels(edges): + return {(u, v): g[u][v]["weight"] for u, v in edges} + + if straight: + nx.draw_networkx_edges(g, pos, edgelist=straight, node_size=node_size) + nx.draw_networkx_edge_labels(g, pos, edge_labels=_edge_labels(straight)) + if curved: + curved = list(curved) + nx.draw_networkx_edges( + g, pos, edgelist=curved, node_size=node_size, connectionstyle=connectionstyle + ) + nx.draw_networkx_edge_labels( + g, pos, edge_labels=_edge_labels(curved), connectionstyle=connectionstyle + ) plt.show() @@ -88,12 +128,12 @@ def spy(M, *, centered=False, show=True, figure=None, axes=None, figsize=None, * plt.show() if axes is None: if figure is None: - fig = mpl.figure.Figure(figsize=figsize) - axes = fig.subplots() + figure = mpl.figure.Figure(figsize=figsize) + axes = figure.subplots() if kwargs.get("markersize") is None: # Make the square markers "fill" their space markersize = min(axes.bbox.width / A.shape[1], axes.bbox.height / A.shape[0]) - kwargs["markersize"] = max(0.002, markersize * 72 / fig.dpi) + kwargs["markersize"] = max(0.002, markersize * 72 / axes.figure.dpi) axes.spy(A, **kwargs) # Fix offsets if not centered: @@ -103,6 +143,38 @@ def spy(M, *, centered=False, show=True, figure=None, axes=None, figsize=None, * return axes.figure +def _matrix_to_dataframe(M): + """Build the ``(row, col, val)`` DataFrame that ``datashade`` rasterizes. + + Factored out of ``datashade`` so the coordinate convention can be checked + without rendering an interactive plot (see ``_cell_centered_limits``). + """ + np, pd = _get_imports(["np", "pd"], "datashade") + rows, cols, vals = M.to_coo() + max_int = np.iinfo(np.int64).max + if M.nrows > max_int and rows.max() > max_int: + rows = rows.astype(np.float64) + else: + rows = rows.astype(np.int64) + if M.ncols > max_int and cols.max() > max_int: + cols = cols.astype(np.float64) + else: + cols = cols.astype(np.int64) + return pd.DataFrame({"row": rows, "col": cols, "val": vals}) + + +def _cell_centered_limits(M): + """Axis limits that center each element on its integer index, like ``spy``. + + datashader bins points into pixels by ``x_range``/``y_range``. With limits + ``(0, N)`` the pixel for index ``k`` spans ``[k, k+1)``, so an element lands + half a cell to the lower-right of the tick labeled ``k``. Offsetting the + limits by half a cell makes the pixel for index ``k`` span ``[k-0.5, k+0.5)``, + centered on tick ``k`` and matching what ``spy`` draws (python-graphblas #473). + """ + return (-0.5, M.ncols - 0.5), (-0.5, M.nrows - 0.5) + + def datashade(M, agg="count", *, width=None, height=None, opts_kwargs=None, **kwargs): """Interactive plot of the sparsity pattern of a Matrix using hvplot and datashader. @@ -132,19 +204,9 @@ def datashade(M, agg="count", *, width=None, height=None, opts_kwargs=None, **kw spy """ - np, pd, bk, hv, _hp, _ds = _get_imports(["np", "pd", "bk", "hv", "hp", "ds"], "datashade") + bk, hv, _hp, _ds = _get_imports(["bk", "hv", "hp", "ds"], "datashade") if "df" not in kwargs: - rows, cols, vals = M.to_coo() - max_int = np.iinfo(np.int64).max - if M.nrows > max_int and rows.max() > max_int: - rows = rows.astype(np.float64) - else: - rows = rows.astype(np.int64) - if M.ncols > max_int and cols.max() > max_int: - cols = cols.astype(np.float64) - else: - cols = cols.astype(np.int64) - df = pd.DataFrame({"row": rows, "col": cols, "val": vals}) + df = _matrix_to_dataframe(M) else: df = kwargs.pop("df") @@ -183,6 +245,7 @@ def datashade(M, agg="count", *, width=None, height=None, opts_kwargs=None, **kw images.extend(image_row) return hv.Layout(images).cols(ncols) + xlim, ylim = _cell_centered_limits(M) kwds = { "x": "col", "y": "row", @@ -192,8 +255,8 @@ def datashade(M, agg="count", *, width=None, height=None, opts_kwargs=None, **kw "frame_height": height, "cmap": "fire", "cnorm": "eq_hist", - "xlim": (0, M.ncols), - "ylim": (0, M.nrows), + "xlim": xlim, + "ylim": ylim, "rasterize": True, "flip_yaxis": True, "hover": True, diff --git a/pyproject.toml b/pyproject.toml index 9fe272c98..c8f4f9ab3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -212,9 +212,6 @@ filterwarnings = [ [tool.coverage.run] branch = true source = ["graphblas"] -omit = [ - "graphblas/viz.py", # TODO: test and get coverage for viz.py -] [tool.coverage.report] ignore_errors = false @@ -225,7 +222,7 @@ skip_empty = true exclude_lines = ["pragma: no cover", "raise AssertionError", "raise NotImplementedError"] [tool.codespell] -ignore-words-list = "coo,ba" +ignore-words-list = "bu,coo,ba" [tool.ruff] # https://github.com/charliermarsh/ruff/ @@ -386,6 +383,7 @@ ignore = [ "graphblas/core/operator/base.py" = ["S102"] # exec is used for UDF "graphblas/core/operator/udt_utils.py" = ["S102"] # exec is used for UDT op codegen "graphblas/monoid/numpy.py" = ["PLW0108"] # lambda is needed for numba.njit +"benchmarks/*.py" = ["T201", "B015", "B018"] # asv: bare exprs are the measurement; _verify prints "graphblas/core/ss/matrix.py" = [ "NPY002", # numba doesn't support rng generator yet "PLR1730", diff --git a/scripts/autogenerate.py b/scripts/autogenerate.py index 5f1116674..95a50fe94 100755 --- a/scripts/autogenerate.py +++ b/scripts/autogenerate.py @@ -11,8 +11,37 @@ Modifying infix-methods is much less common, but should be run if you want to modify it. +Pass --check to verify the generated files on disk match a fresh regeneration without +modifying anything. This exits nonzero (and names the drifted files) when they differ, +which is what the drift test uses to catch a hand-edited or stale generated block. + """ +import sys +from pathlib import Path + +# For a script, sys.path[0] is the script's own directory rather than the caller's cwd, so +# a bare `import graphblas` resolves to whatever is installed. That coincides with this +# checkout in an ordinary dev setup and diverges in a worktree, where the generator would +# then rewrite, or validate, a tree nobody asked about while looking perfectly normal. +_REPO_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(_REPO_ROOT)) + +# --check prefixes the package it validated with this, so the drift test can assert it +# checked the checkout it lives in rather than infer correctness from an exit code. +_PACKAGE_LINE = "graphblas package: " + +# Files whose auto-generated blocks the generators (re)write. During --check, reads and +# writes are redirected to scratch copies of these so the real tree is never touched. +_GENERATED_FILES = ( + "automethods.py", + "infixmethods.py", + "scalar.py", + "vector.py", + "matrix.py", + "infix.py", +) + def main(): from graphblas.core.automethods import _main as auto_main @@ -22,5 +51,62 @@ def main(): infix_main() +def _parsed(path): + """Source parsed to an AST dump: identical for two files that differ only in layout.""" + import ast + + return ast.dump(ast.parse(path.read_bytes())) + + +def check(): + """Regenerate into a scratch tree and report any drift. + + Returns 0 when every generated file matches a fresh regeneration, else 1. + + The comparison is on parsed syntax rather than bytes, which keeps the check + independent of `black`. The generators shell out to black when it is on PATH and skip + it when it is not, so a byte comparison reports drift on automethods.py and + infixmethods.py in every environment lacking it, and black is not a test dependency. + Layout is already enforced repo-wide by black in pre-commit and the lint job, so the + drift left for this check to catch is a generated block whose content is stale. + """ + import shutil + import tempfile + from pathlib import Path + + import graphblas + from graphblas.core import automethods, infixmethods + + # Report the tree actually validated. For a script sys.path[0] is the script's own + # directory, so without the sys.path fix above this silently checks whichever + # graphblas is installed; the drift test asserts on this line. + print(f"{_PACKAGE_LINE}{Path(graphblas.__file__).resolve().parent}") + + src_dir = Path(automethods.__file__).parent + with tempfile.TemporaryDirectory(prefix="autogen_check_") as tmp: + scratch = Path(tmp) + for name in _GENERATED_FILES: + shutil.copyfile(src_dir / name, scratch / name) + automethods._main(base_dir=scratch, callblack=False) + infixmethods._main(base_dir=scratch, callblack=False) + + drifted = [ + name for name in _GENERATED_FILES if _parsed(scratch / name) != _parsed(src_dir / name) + ] + + if drifted: + print("Auto-generated code is out of date; run `python scripts/autogenerate.py`.") + print("Drifted file(s):") + for name in drifted: + print(f" graphblas/core/{name}") + return 1 + print("Auto-generated code is up to date.") + return 0 + + if __name__ == "__main__": + import sys + + if "--check" in sys.argv[1:]: + sys.exit(check()) main()