diff --git a/python/python/tests/test_expr.py b/python/python/tests/test_expr.py index 0f49231f1..6cfc09922 100644 --- a/python/python/tests/test_expr.py +++ b/python/python/tests/test_expr.py @@ -80,6 +80,67 @@ class TestExprConstruction: e = lit(d) assert isinstance(e, Expr) + @pytest.mark.parametrize( + "value, decimal_type", + [ + ("1E-7", pa.decimal128(12, 7)), + ("-1.23E-7", pa.decimal128(12, 9)), + ("1E+3", pa.decimal128(12, 0)), + ("-1.23E+5", pa.decimal128(12, 0)), + ("0E-7", pa.decimal128(12, 7)), + ("-0E+3", pa.decimal128(12, 0)), + ("1E-38", pa.decimal128(38, 38)), + ("1E+37", pa.decimal128(38, 0)), + ("1.234567890123456789012345678901234567E+37", pa.decimal128(38, 0)), + ("0E+128", pa.decimal128(1, 0)), + ], + ) + def test_lit_decimal_exponent_query(self, tmp_path, value, decimal_type): + expected = Decimal(value) + table = lancedb.connect(tmp_path).create_table( + "decimals", + pa.table( + { + "id": [1, 2], + # Arrow's Python ingestion rejects 0E+128 even though its + # numeric value is representable. Store canonical zero to + # exercise the expression literal independently of ingestion. + "price": pa.array( + [Decimal(0) if expected.is_zero() else expected, None], + type=decimal_type, + ), + } + ), + ) + predicate = col("price") == lit(expected) + result = table.search().where(predicate).to_arrow() + assert result["id"].to_pylist() == [1] + assert result["price"].to_pylist() == [expected] + # Exercise the SQL serializer and parser as well as the typed query path. + sql_result = table.search().where(predicate.to_sql()).to_arrow() + assert sql_result["id"].to_pylist() == [1] + assert sql_result["price"].to_pylist() == [expected] + + @pytest.mark.parametrize( + "value", + [ + "NaN", + "sNaN", + "Infinity", + "-Infinity", + "1E-39", + "1E+38", + "1E+129", + "1E+1000000000", + "0E+129", + "12345678901234567890123456789012345678E+1", + "9" * 39, + ], + ) + def test_lit_decimal_unrepresentable_raises(self, value): + with pytest.raises(ValueError): + lit(Decimal(value)) + class TestExprOperators: def test_eq_operator(self): diff --git a/python/src/expr.rs b/python/src/expr.rs index 79b448fdf..bb779bc1d 100644 --- a/python/src/expr.rs +++ b/python/src/expr.rs @@ -240,22 +240,60 @@ pub fn expr_lit(value: Bound<'_, PyAny>) -> PyResult { fn parse_decimal(s: &str) -> PyResult<(i128, u8, i8)> { let s = s.trim(); - let dot_pos = s.find('.'); - let scale = if let Some(pos) = dot_pos { - (s.len() - pos - 1) as i8 - } else { - 0 + let (mantissa, exponent) = match s.split_once(['e', 'E']) { + Some((mantissa, exponent)) => ( + mantissa, + exponent.parse::().map_err(|e| { + PyValueError::new_err(format!("failed to parse decimal exponent: {}", e)) + })?, + ), + None => (s, 0), }; + let fractional_digits = mantissa.find('.').map_or(0, |pos| mantissa.len() - pos - 1); + // Keep the coefficient intact instead of expanding an exponent into zeros. + // Arrow Decimal128 has at most 38 digits and an i8 scale (at most 38). + let mut scale = i64::try_from(fractional_digits) + .ok() + .and_then(|fractional_digits| fractional_digits.checked_sub(exponent)) + .and_then(|scale| i8::try_from(scale).ok()) + .filter(|scale| *scale <= 38) + .ok_or_else(|| PyValueError::new_err("decimal scale is outside Decimal128 range"))?; - let digits = s.replace('.', ""); - let val = digits + let digits = mantissa.replace('.', ""); + let mut val = digits .parse::() .map_err(|e| PyValueError::new_err(format!("failed to parse decimal digits: {}", e)))?; - // Precision is total number of digits - let precision = digits.trim_start_matches('-').len() as u8; + // Positive scale cannot exceed precision, including for tiny values such as 1E-38. + let mut precision = digits + .trim_start_matches('-') + .len() + .max(usize::try_from(scale).unwrap_or(0)); + if scale < 0 { + // Arrow cannot cast SQL text to a decimal with negative scale. Normalize + // positive exponents within Decimal128's bounded coefficient range. + if val != 0 { + let shift = -i16::from(scale) as u32; + precision += shift as usize; + if precision > 38 { + return Err(PyValueError::new_err( + "decimal precision exceeds Decimal128 maximum of 38", + )); + } + val = 10_i128 + .checked_pow(shift) + .and_then(|factor| val.checked_mul(factor)) + .ok_or_else(|| PyValueError::new_err("decimal coefficient exceeds i128 range"))?; + } + scale = 0; + } + if precision > 38 { + return Err(PyValueError::new_err( + "decimal precision exceeds Decimal128 maximum of 38", + )); + } - Ok((val, precision, scale)) + Ok((val, precision as u8, scale)) } /// Call an arbitrary registered SQL function by name.