fix: floor // on DuckDB for negative operands and floats

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":

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"]