diff --git a/src/narwhals/_duckdb/utils.py b/src/narwhals/_duckdb/utils.py index ae84fe1d99..ec7669d50e 100644 --- a/src/narwhals/_duckdb/utils.py +++ b/src/narwhals/_duckdb/utils.py @@ -11,7 +11,7 @@ from narwhals.exceptions import ColumnNotFoundError if TYPE_CHECKING: - from collections.abc import Mapping, Sequence + from collections.abc import Callable, Mapping, Sequence from typing import TypeAlias from duckdb import DuckDBPyRelation @@ -400,12 +400,28 @@ def catch_duckdb_exception( return exception -_BINARY_OPS = {"floordiv": operator.floordiv, "and": operator.and_} +def _floordiv(dividend: Expression, divisor: Expression) -> Expression: + # DuckDB's `//` truncates integers and is true division on floats. + quotient = dividend // divisor + # `floor` would turn integers into DOUBLE. + truncated = F("trunc", quotient) + # `typeof` folds at bind time, so DuckDB only evaluates the arm for this dtype. + is_float = F("typeof", quotient).isin(lit("FLOAT"), lit("DOUBLE")) + is_negative_inexact = when(is_float, quotient < truncated).otherwise( + (dividend % divisor != lit(0)) & ((dividend < lit(0)) != (divisor < lit(0))) + ) + return when(is_negative_inexact, truncated - lit(1)).otherwise(truncated) + + +_BINARY_OPS: Mapping[str, Callable[[Expression, Expression], Expression]] = { + "floordiv": _floordiv, + "and": operator.and_, +} def function(name: str, *args: Expression) -> Expression: if name in _BINARY_OPS: - return _BINARY_OPS[name](*args) # type: ignore[no-any-return] + return _BINARY_OPS[name](*args) if name in {"isnull", "isnotnull"}: return getattr(args[0], name)() # type: ignore[no-any-return] if name == "count_distinct": diff --git a/tests/expr_and_series/arithmetic_test.py b/tests/expr_and_series/arithmetic_test.py index 4deab5fde6..add0938110 100644 --- a/tests/expr_and_series/arithmetic_test.py +++ b/tests/expr_and_series/arithmetic_test.py @@ -11,6 +11,7 @@ DASK_VERSION, DUCKDB_VERSION, PANDAS_VERSION, + POLARS_VERSION, Constructor, ConstructorEager, assert_equal_data, @@ -37,8 +38,6 @@ def test_arithmetic_expr( constructor: Constructor, request: pytest.FixtureRequest, ) -> None: - if "duckdb" in str(constructor) and attr == "__floordiv__": - request.applymarker(pytest.mark.xfail) if attr == "__mod__" and any( x in str(constructor) for x in ["pandas_pyarrow", "modin_pyarrow"] ): @@ -123,6 +122,97 @@ def test_rmod_negative_operands( assert_equal_data(result, {"literal": expected}) +@pytest.mark.parametrize( + ("data", "divisor", "expected"), + [([-7, 7, -1, None], 2, [-4, 3, -1, None]), ([7, -7], -2, [-4, 3])], +) +def test_floordiv_negative_operands( + data: list[Any], divisor: int, expected: list[Any], constructor: Constructor +) -> None: + df = nw.from_native(constructor({"a": data})) + result = df.select(nw.col("a") // divisor) + assert_equal_data(result, {"a": expected}) + + +@pytest.mark.parametrize( + ("data", "dividend", "expected"), + [([2, 3, None], -7, [-4, -3, None]), ([-2, 2], 7, [-4, 3])], +) +def test_rfloordiv_negative_operands( + data: list[Any], dividend: int, expected: list[Any], constructor: Constructor +) -> None: + df = nw.from_native(constructor({"a": data})) + result = df.select(dividend // nw.col("a")) + assert_equal_data(result, {"literal": expected}) + + +def test_floordiv_large_integers( + constructor: Constructor, request: pytest.FixtureRequest +) -> None: + if "polars" in str(constructor) and POLARS_VERSION < (0, 20, 7): + pytest.skip(reason="floors through float64") + if any(x in str(constructor) for x in ("ibis", "sqlframe", "pyspark")): + request.applymarker(pytest.mark.xfail(reason="floors through float64")) + # Past 2**53, so flooring a float quotient is off by one. + big = 2**62 + 1 + i64_max = 2**63 - 1 + data = {"a": [big, -big, 10, -(2**63)], "b": [3, 3, i64_max, 3]} + df = nw.from_native(constructor(data)) + result = df.select(nw.col("a") // nw.col("b"), i64_max // nw.col("a")) + expected = { + "a": [1537228672809129301, -1537228672809129302, 0, -3074457345618258603], + "literal": [1, -2, 922337203685477580, -1], + } + assert_equal_data(result, expected) + + +@pytest.mark.parametrize("dtype", [nw.UInt32, nw.UInt64]) +def test_floordiv_unsigned_dtype( + dtype: type[nw.UInt32 | nw.UInt64], + constructor: Constructor, + request: pytest.FixtureRequest, +) -> None: + if any(x in str(constructor) for x in ["sqlframe", "pyspark"]): + pytest.skip(reason="Spark has no unsigned integers") + if any(x in str(constructor) for x in ["pandas_pyarrow", "modin_pyarrow", "ibis"]): + request.applymarker(pytest.mark.xfail(reason="returns Int64")) + df = nw.from_native(constructor({"a": [7, 9], "b": [2, 4]})) + df = df.select(nw.all().cast(dtype)) + result = df.select(nw.col("a") // nw.col("b"), nw.lit(7, dtype) // nw.col("b")) + assert result.collect_schema() == {"a": dtype, "literal": dtype} + assert_equal_data(result, {"a": [3, 2], "literal": [3, 1]}) + + +@pytest.mark.parametrize( + ("left", "right", "expected"), + [ + (7.0, 2.0, 3.0), + (-7.0, 2.0, -4.0), + (1.0, float("inf"), 0.0), + (1.0, 0.1, 10.0), + (-1.0, 0.1, -10.0), + ], +) +def test_floordiv_floats( + left: float, + right: float, + expected: float, + constructor: Constructor, + request: pytest.FixtureRequest, +) -> None: + if ( + left // right != expected + and "pyarrow" not in str(constructor) + and any(x in str(constructor) for x in ("pandas", "modin", "cudf", "dask")) + ): + request.applymarker( + pytest.mark.xfail(reason="pandas follows Python, where 1.0 // 0.1 is 9.0") + ) + df = nw.from_native(constructor({"a": [left], "b": [right]})) + result = df.select(nw.col("a") // nw.col("b"), left // nw.col("b")) + assert_equal_data(result, {"a": [expected], "literal": [expected]}) + + @pytest.mark.parametrize( ("attr", "rhs", "expected"), [ @@ -254,9 +344,7 @@ def test_arithmetic_expr_left_literal( constructor: Constructor, request: pytest.FixtureRequest, ) -> None: - if ("duckdb" in str(constructor) and attr == "__floordiv__") or ( - "dask" in str(constructor) and DASK_VERSION < (2024, 10) - ): + if "dask" in str(constructor) and DASK_VERSION < (2024, 10): request.applymarker(pytest.mark.xfail) if attr == "__mod__" and any( x in str(constructor) for x in ["pandas_pyarrow", "modin_pyarrow"]