diff --git a/src/blosc2/lazyexpr.py b/src/blosc2/lazyexpr.py index 2229cf7f2..5e44ed362 100644 --- a/src/blosc2/lazyexpr.py +++ b/src/blosc2/lazyexpr.py @@ -358,6 +358,23 @@ def _find_constructor_call(expression: str, constructor: str) -> re.Match | None "//": operator.floordiv, "%": operator.mod, } +# NumPy operations of the numeric operators and functions whose result dtype is not just the +# promoted operand dtype: true division and the float functions turn integers (and bools) into +# floats, and //, %, ** and the shifts, which NumPy has no bool loops for, turn bools into int8 +_result_dtype_ufuncs = { + "/": operator.truediv, + "//": operator.floordiv, + "%": operator.mod, + "**": operator.pow, # not np.power: NumPy squares for x ** 2, which keeps bools int8 + "<<": operator.lshift, + ">>": operator.rshift, + "arctan2": np.arctan2, + "copysign": np.copysign, + "hypot": np.hypot, + "nextafter": np.nextafter, + "pow": np.power, + "power": np.power, +} logical_ops = ["&", "|", "^", "~"] not_complex_ops = ["maximum", "minimum", "<", "<=", ">", ">="] funcs_2args = ( @@ -3624,12 +3641,24 @@ def fuse_expressions(expr, new_base, dup_op): return new_expr +def _operand_dtype(value): + """blosc2.result_type(value), without building a NumPy array for an array operand.""" + dtype = None if np.isscalar(value) else getattr(value, "dtype", None) + if dtype is not None: + dtype = convert_dtype(dtype) + if dtype.kind != "V": + # np.result_type() gives the native byte order, which the dtype checks below compare against + return np.result_type(dtype) + # Structured and void dtypes too: blosc2.result_type() rejects some of them + return blosc2.result_type(value) + + def check_dtype(op, value1, value2): if op in ("contains", "startswith", "endswith"): return np.dtype(np.bool_) - v1_dtype = blosc2.result_type(value1) - v2_dtype = v1_dtype if value2 is None else blosc2.result_type(value2) + v1_dtype = _operand_dtype(value1) + v2_dtype = v1_dtype if value2 is None else _operand_dtype(value2) if op in not_complex_ops and (v1_dtype == np.complex128 or v2_dtype == np.complex128): # Ensure that throw exception for functions which don't support complex args raise ValueError(f"Invalid operand type for {op}: {v1_dtype, v2_dtype}") @@ -3642,13 +3671,23 @@ def check_dtype(op, value1, value2): if v2_dtype not in (np.bool_, np.int32, np.int64): raise ValueError(f"Invalid operand type for {op}: {v2_dtype}") - if op == "/": - if v1_dtype == np.int32 and v2_dtype == np.int32: - return blosc2.float32 - if np.issubdtype(v1_dtype, np.integer) and np.issubdtype(v2_dtype, np.integer): - return blosc2.float64 + if op == "/" and v1_dtype == np.int32 and v2_dtype == np.int32: + return blosc2.float32 + + ufunc = _result_dtype_ufuncs.get(op) + if ufunc is not None and value2 is not None and v1_dtype.kind in "biufc" and v2_dtype.kind in "biufc": + # Ask NumPy on empty operands, but with a scalar next to an array as it is, so that NumPy + # promotes it as usual; what it rejects is left to the evaluation + scalars = np.isscalar(value1) and np.isscalar(value2) + probes = [ + v if np.isscalar(v) and not scalars else np.empty(0, dt) + for v, dt in ((value1, v1_dtype), (value2, v2_dtype)) + ] + with contextlib.suppress(TypeError, ValueError, OverflowError): + return ufunc(*probes).dtype if op == "arctan2": + # Except for arctan2, which rejects what NumPy rejects (complex, datetimes, strings) here return np.arctan2(np.empty(0, v1_dtype), np.empty(0, v2_dtype)).dtype if (v1_dtype.kind in "Mm" or v2_dtype.kind in "Mm") and op in _datetime_binops: diff --git a/tests/ndarray/test_lazyexpr.py b/tests/ndarray/test_lazyexpr.py index 4be24af5d..324cd71b4 100644 --- a/tests/ndarray/test_lazyexpr.py +++ b/tests/ndarray/test_lazyexpr.py @@ -45,6 +45,125 @@ def test_arctan2_preserves_fractional_results(dtype, fused): np.testing.assert_allclose(blosc2.lazyexpr("arctan2(a, b)")[:], expected, rtol=1e-6) +_RESULT_DTYPE_EXPRESSIONS = [ + "hypot({n}, b)", + "copysign({n}, b)", + "nextafter({n}, b)", + "arctan2({n}, b)", + "hypot(a, b * 1)", + "hypot({n}, 2)", + "arctan2({n}, 2.5)", + "(a < b) / b", + "a / 2", + "a / b", + "{n} / b", + "{n} // b", + "{n} % b", + "a ** b", + "{n} ** 2", + "sin({n})", + "sqrt(a)", +] + + +# Known wrong values on WebAssembly: there compute() and slices also run the bundled miniexpr, which +# computes these in integer arithmetic (0.0 for NumPy's 0.5). strict, so that a fix shows up as XPASS +_MINIEXPR_TRUE_DIVISION = ( + "miniexpr computes integer true division in integer arithmetic: " + "https://github.com/Blosc/miniexpr/issues/26" +) +_MINIEXPR_FLOAT_FUNCTIONS = ( + "miniexpr computes hypot() and nextafter() of integers in integer arithmetic: " + "https://github.com/Blosc/miniexpr/issues/27" +) +_WASM_KNOWN_FAILURES = { + "hypot({n}, b)": ((np.uint16, np.int32, np.int64), _MINIEXPR_FLOAT_FUNCTIONS), + "nextafter({n}, b)": ((np.int32, np.int64), _MINIEXPR_FLOAT_FUNCTIONS), + "(a < b) / b": ((np.int8, np.uint16, np.int32, np.int64), _MINIEXPR_TRUE_DIVISION), + "a / b": ((np.int8, np.uint16, np.int64), _MINIEXPR_TRUE_DIVISION), + "{n} / b": ((np.int8, np.uint16, np.int64), _MINIEXPR_TRUE_DIVISION), +} + + +@pytest.mark.parametrize("dtype", [np.bool_, np.int8, np.uint16, np.int32, np.int64]) +@pytest.mark.parametrize("expression", _RESULT_DTYPE_EXPRESSIONS) +@pytest.mark.parametrize("build", ["functions", "string"]) +def test_lazyexpr_result_dtype_matches_numpy(dtype, expression, build, request): + # Float functions and true division of integer (or bool) subexpressions were given an integer + # (or bool) result dtype, and bool //, % and ** a bool one, so their results were truncated + if dtype is np.int32 and expression in ("a / b", "{n} / b"): + pytest.skip("int32 / int32 gives float32 on purpose") + if "nextafter" in expression and dtype not in (np.int32, np.int64): + pytest.skip("a float16 or float32 nextafter() is computed in float64 on every path") + known_dtypes, reason = _WASM_KNOWN_FAILURES.get(expression, ((), None)) + if blosc2.IS_WASM and dtype in known_dtypes: + request.applymarker(pytest.mark.xfail(reason=reason, strict=True)) + # A subexpression of the operand dtype (bool + bool is an or in NumPy) + expression = expression.format(n="(a & b)" if dtype is np.bool_ else "(a + b)") + na = (np.arange(1200).reshape(30, 40) % 5).astype(dtype) + nb = (np.arange(1200).reshape(30, 40) % 3 + 1).astype(dtype) + a = blosc2.asarray(na, chunks=(10, 40), blocks=(5, 40)) + b = blosc2.asarray(nb, chunks=(10, 40), blocks=(5, 40)) + namespace = {"hypot": np.hypot, "copysign": np.copysign, "nextafter": np.nextafter} + namespace.update(arctan2=np.arctan2, sin=np.sin, sqrt=np.sqrt) + with np.errstate(all="ignore"): + expected = eval(expression, namespace, {"a": na, "b": nb}) + if dtype is np.int32 and expression == "a / 2" and blosc2.result_type(2) == np.int32: + # A Python int is int32 here (as on wasm32), so this is int32 / int32, which gives float32 + # on purpose; elsewhere it is int32 / int64, float64 as in NumPy + expected = expected.astype(np.float32) + if build == "string": + expr = blosc2.lazyexpr(expression, {"a": a, "b": b}) + else: + expr = eval(expression, {name: getattr(blosc2, name) for name in namespace}, {"a": a, "b": b}) + assert expr.dtype == expected.dtype + for res in (expr[:], expr[...], expr.compute()[:], expr[:7]): + assert res.dtype == expected.dtype + # Values on the paths whose typing follows NumPy here; expr[:] and reductions of integer + # true division and two-argument float functions still compute in integer arithmetic in + # the bundled miniexpr + rtol = 1e-3 if expected.dtype == np.float16 else 1e-6 + np.testing.assert_allclose(expr.compute()[:], expected, rtol=rtol) + np.testing.assert_allclose(expr[:7], expected[:7], rtol=rtol) + + +@pytest.mark.parametrize("dtype", [np.int32, np.int64, np.float64]) +@pytest.mark.parametrize("op", ["&", "|", "^", "<<", ">>", "+", "*", "/", "//", "%", "**", "<", "hypot"]) +def test_lazyexpr_non_native_byte_order(dtype, op): + # The operand dtype checks compare against native dtypes: a non-native operand must give the + # same result dtype as a native one (no "Invalid operand type for &: >i4", and >i4 / >i4 + # keeps the float32 of int32 / int32) + if dtype is np.float64 and op in ("&", "|", "^", "<<", ">>"): + pytest.skip("no float bitwise operators") + native = np.dtype(dtype) + swapped = native.newbyteorder() + na = np.arange(1, 7, dtype=native) + nb = np.array([1, 2, 1, 3, 1, 2], dtype=native) + + def build(a, b): + return blosc2.hypot(a, b) if op == "hypot" else eval(f"a {op} b") + + reference = build(blosc2.asarray(na), nb) + if op == "/" and dtype is np.int32: + assert np.dtype(reference.dtype) == np.float32 + expr = build(blosc2.asarray(na), nb.astype(swapped)) + assert expr.dtype == reference.dtype + np.testing.assert_array_equal(expr[:], reference[:]) + np.testing.assert_array_equal(expr.compute()[:], reference.compute()[:]) + # Only the dtype for non-native NDArray operands: their values are a separate problem + sa, sb = na.astype(swapped), nb.astype(swapped) + for a, b in ((na, sb), (sa, nb), (sa, sb)): + assert build(blosc2.asarray(a), blosc2.asarray(b)).dtype == reference.dtype + + +@pytest.mark.parametrize("other", [np.ones(4, dtype=np.int32), 2, 2.5]) +def test_arctan2_rejects_complex_when_built(other): + c = blosc2.asarray(np.ones(4, dtype=np.complex128)) + for x in (c, c + c): + with pytest.raises(TypeError): + blosc2.arctan2(x, other) + + _UNARY_FUNCTIONS = [ "sin", "cos",