Skip to content

Commit 9a6131a

Browse files
committed
Merge branch 'main' into normalize-spaces
2 parents 2d19d9f + 0d97293 commit 9a6131a

2 files changed

Lines changed: 79 additions & 10 deletions

File tree

‎python/cudf_polars/cudf_polars/dsl/expressions/binaryop.py‎

Lines changed: 28 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66

77
from __future__ import annotations
88

9-
from decimal import Decimal
9+
from decimal import Decimal, localcontext
1010
from typing import TYPE_CHECKING, ClassVar
1111

1212
import polars as pl
@@ -17,6 +17,7 @@
1717

1818
from cudf_polars.containers import Column
1919
from cudf_polars.dsl.expressions.base import ExecutionContext, Expr
20+
from cudf_polars.dsl.expressions.literal import Literal
2021

2122
pl_expr = polars._expr_nodes
2223

@@ -30,18 +31,37 @@
3031
_DECIMAL128_MAX_PRECISION = 38
3132

3233

33-
def _decimal_rescale_may_overflow(left: DataType, right: DataType) -> bool:
34+
def _decimal_precision(node: Expr) -> int:
35+
"""
36+
The number of digits needed for the values of a Decimal expression.
37+
38+
A literal only needs as many digits as its value, which can be far fewer
39+
than its type declares: polars gives the literal ``24`` the type
40+
``Decimal(38, 0)``.
41+
"""
42+
dtype = node.dtype.polars_type
43+
assert isinstance(dtype, pl.Decimal)
44+
if isinstance(node, Literal) and isinstance(node.value, Decimal):
45+
# Polars writes a Decimal literal with exactly the scale of its dtype, so
46+
# the scaled value is a whole number. The default context rounds to 28
47+
# digits, which would change the count near the 38 digit limit.
48+
with localcontext(prec=2 * _DECIMAL128_MAX_PRECISION):
49+
return max(len(str(abs(int(node.value.scaleb(dtype.scale))))), 1)
50+
return dtype.precision or _DECIMAL128_MAX_PRECISION
51+
52+
53+
def _decimal_rescale_may_overflow(left: Expr, right: Expr) -> bool:
3454
"""Whether rescaling two Decimals to their finer scale can exceed 38 digits."""
3555
if not (
36-
plc.traits.is_fixed_point(left.plc_type)
37-
and plc.traits.is_fixed_point(right.plc_type)
56+
plc.traits.is_fixed_point(left.dtype.plc_type)
57+
and plc.traits.is_fixed_point(right.dtype.plc_type)
3858
):
3959
return False
40-
lt, rt = left.polars_type, right.polars_type
60+
lt, rt = left.dtype.polars_type, right.dtype.polars_type
4161
assert isinstance(lt, pl.Decimal)
4262
assert isinstance(rt, pl.Decimal)
43-
lp = lt.precision or _DECIMAL128_MAX_PRECISION
44-
rp = rt.precision or _DECIMAL128_MAX_PRECISION
63+
lp = _decimal_precision(left)
64+
rp = _decimal_precision(right)
4565
scale = max(lt.scale, rt.scale)
4666
return max(lp - lt.scale, rp - rt.scale) + scale > _DECIMAL128_MAX_PRECISION
4767

@@ -66,9 +86,7 @@ def __init__(
6686
self.op = op
6787
self.children = (left, right)
6888
self.is_pointwise = True
69-
if op in BinOp._COMPARISONS and _decimal_rescale_may_overflow(
70-
left.dtype, right.dtype
71-
):
89+
if op in BinOp._COMPARISONS and _decimal_rescale_may_overflow(left, right):
7290
# Polars 2.1 compares decimals of different scales exactly, but
7391
# libcudf rescales to the finer scale in 128 bits and wraps.
7492
raise NotImplementedError(

‎python/cudf_polars/tests/expressions/test_numeric_binops.py‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -332,3 +332,54 @@ def test_decimal_comparison_mixed_scale_overflow(
332332
assert_ir_translation_raises(q, engine, NotImplementedError)
333333
else:
334334
assert_gpu_result_equal(q, engine=engine)
335+
336+
337+
@pytest.mark.parametrize("dtype", [pl.Decimal(15, 2), pl.Decimal(38, 2)])
338+
@pytest.mark.parametrize("literal", [0, 24, -6, 10_000])
339+
def test_decimal_comparison_with_integer_literal(engine: pl.GPUEngine, dtype, literal):
340+
# polars types an integer literal as Decimal(38, 0), but its value needs
341+
# far fewer digits than that, so rescaling it cannot overflow.
342+
biggest = "9" * (dtype.precision - dtype.scale) + "." + "9" * dtype.scale
343+
df = pl.DataFrame({"a": ["-2.5", "0", "24", str(literal), biggest, None]}).cast(
344+
dtype
345+
)
346+
a = pl.col("a")
347+
q = df.lazy().select(
348+
lt=a < literal,
349+
le=a <= literal,
350+
gt=a > literal,
351+
ge=a >= literal,
352+
eq=a == literal,
353+
ne=a != literal,
354+
)
355+
assert_gpu_result_equal(q, engine=engine)
356+
357+
358+
def test_decimal_comparison_of_expression_with_integer_literal(
359+
engine: pl.GPUEngine,
360+
):
361+
df = pl.DataFrame({"a": ["1.50", "2.25", "3.00"]}).cast(pl.Decimal(7, 2)).lazy()
362+
q = df.select(gt=(pl.col("a") * 2).sum() > 0, lt=pl.col("a") * 2 < 5)
363+
assert_gpu_result_equal(q, engine=engine)
364+
365+
366+
def test_decimal_comparison_with_literal_that_may_overflow(engine: pl.GPUEngine):
367+
# The literal has a finer scale than the column, so rescaling the column
368+
# to the literal's scale needs 36 + 3 digits.
369+
df = pl.DataFrame({"a": ["1.00", "2.00"]}).cast(pl.Decimal(38, 2)).lazy()
370+
q = df.select(gt=pl.col("a") > pl.lit(Decimal("1.005")))
371+
assert_ir_translation_raises(q, engine, NotImplementedError)
372+
373+
374+
def test_decimal_comparison_with_literal_of_maximum_precision(engine: pl.GPUEngine):
375+
# Neither operand needs rescaling, so a literal of 38 digits is supported.
376+
df = pl.DataFrame({"a": ["1", "9" * 38]}).cast(pl.Decimal(38, 0)).lazy()
377+
q = df.select(lt=pl.col("a") < pl.lit(Decimal("9" * 38)))
378+
assert_gpu_result_equal(q, engine=engine)
379+
380+
381+
@pytest.mark.parametrize("literal", [Decimal("27.5"), 27.5])
382+
def test_decimal_comparison_with_fractional_literal(engine: pl.GPUEngine, literal):
383+
df = pl.DataFrame({"a": ["1.00", "27.50", "30.00", None]}).cast(pl.Decimal(15, 2))
384+
q = df.lazy().select(lt=pl.col("a") < literal, ge=pl.col("a") >= literal)
385+
assert_gpu_result_equal(q, engine=engine)

0 commit comments

Comments
 (0)