Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 13 additions & 9 deletions graphblas/core/operator/binary.py
Original file line number Diff line number Diff line change
Expand Up @@ -1004,15 +1004,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__(
Expand All @@ -1037,6 +1034,13 @@ def __init__(
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

__call__ = TypedBuiltinBinaryOp.__call__
is_commutative = TypedBuiltinBinaryOp.is_commutative
Expand Down
40 changes: 40 additions & 0 deletions graphblas/tests/test_op.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,8 @@
import itertools
import os
import subprocess
import sys
from pathlib import Path

import numpy as np
import pytest
Expand Down Expand Up @@ -3768,3 +3772,39 @@ def test_operator_namespace_typo_suggestions():
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