66
77from __future__ import annotations
88
9- from decimal import Decimal
9+ from decimal import Decimal , localcontext
1010from typing import TYPE_CHECKING , ClassVar
1111
1212import polars as pl
1717
1818from cudf_polars .containers import Column
1919from cudf_polars .dsl .expressions .base import ExecutionContext , Expr
20+ from cudf_polars .dsl .expressions .literal import Literal
2021
2122pl_expr = polars ._expr_nodes
2223
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 (
0 commit comments