feat: parse_float_as_decimal supports scientific notation and Decimal256
#13806
+245
−47
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -19,10 +19,13 @@ use crate::planner::{ContextProvider, PlannerContext, SqlToRel}; | |
| use arrow::compute::kernels::cast_utils::{ | ||
| parse_interval_month_day_nano_config, IntervalParseConfig, IntervalUnit, | ||
| }; | ||
| use arrow::datatypes::{i256, DECIMAL128_MAX_PRECISION}; | ||
| use arrow_schema::{DataType, DECIMAL256_MAX_PRECISION}; | ||
| use bigdecimal::num_bigint::BigInt; | ||
| use bigdecimal::{BigDecimal, Signed, ToPrimitive}; | ||
| use datafusion_common::{ | ||
| internal_datafusion_err, internal_err, not_impl_err, plan_err, DFSchema, | ||
| DataFusionError, Result, ScalarValue, | ||
| }; | ||
| use datafusion_expr::expr::{BinaryExpr, Placeholder}; | ||
| use datafusion_expr::planner::PlannerResult; | ||
|
|
@@ -31,6 +34,9 @@ use log::debug; | |
| use sqlparser::ast::{BinaryOperator, Expr as SQLExpr, Interval, UnaryOperator, Value}; | ||
| use sqlparser::parser::ParserError::ParserError; | ||
| use std::borrow::Cow; | ||
| use std::cmp::Ordering; | ||
| use std::ops::Neg; | ||
| use std::str::FromStr; | ||
|
|
||
| impl<S: ContextProvider> SqlToRel<'_, S> { | ||
| pub(crate) fn parse_value( | ||
|
|
@@ -84,7 +90,7 @@ impl<S: ContextProvider> SqlToRel<'_, S> { | |
| } | ||
|
|
||
| if self.options.parse_float_as_decimal { | ||
| parse_decimal(unsigned_number, negative) | ||
| } else { | ||
| signed_number.parse::<f64>().map(lit).map_err(|_| { | ||
| DataFusionError::from(ParserError(format!( | ||
|
|
@@ -315,45 +321,84 @@ const fn try_decode_hex_char(c: u8) -> Option<u8> { | |
| } | ||
| } | ||
|
|
||
| /// Returns None if the value can't be converted to i256. | ||
| /// Modified from <https://github.com/apache/arrow-rs/blob/c4dbf0d8af6ca5a19b8b2ea777da3c276807fc5e/arrow-buffer/src/bigint/mod.rs#L303> | ||
| fn bigint_to_i256(v: &BigInt) -> Option<i256> { | ||
| let v_bytes = v.to_signed_bytes_le(); | ||
| match v_bytes.len().cmp(&32) { | ||
| Ordering::Less => { | ||
| let mut bytes = if v.is_negative() { | ||
| [255_u8; 32] | ||
| } else { | ||
| [0; 32] | ||
| }; | ||
| bytes[0..v_bytes.len()].copy_from_slice(&v_bytes[..v_bytes.len()]); | ||
| Some(i256::from_le_bytes(bytes)) | ||
| } | ||
| Ordering::Equal => Some(i256::from_le_bytes(v_bytes.try_into().unwrap())), | ||
| Ordering::Greater => None, | ||
| } | ||
| } | ||
|
|
||
| fn parse_decimal(unsigned_number: &str, negative: bool) -> Result<Expr> { | ||
|
There was a problem hiding this comment. Please add unit tests for this function, to exercise bound checks. There was a problem hiding this comment. I think the new SLT tests have covered them, but please remind me if I've overlooked anything There was a problem hiding this comment. they likely do, but i find it valuable to have unit tests, especially for a code which is well isolated (easy to unit tests) and has so many IFs (valuable to unit test). There was a problem hiding this comment. Make sense. Added in 2077749 |
||
| let mut dec = BigDecimal::from_str(unsigned_number).map_err(|e| { | ||
| DataFusionError::from(ParserError(format!( | ||
| "Cannot parse {unsigned_number} as BigDecimal: {e}" | ||
| ))) | ||
| })?; | ||
| if negative { | ||
| dec = dec.neg(); | ||
| } | ||
|
|
||
| let digits = dec.digits(); | ||
| let (int_val, scale) = dec.into_bigint_and_exponent(); | ||
| if scale < i8::MIN as i64 { | ||
findepi marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| return not_impl_err!( | ||
| "Decimal scale {} exceeds the minimum supported scale: {}", | ||
| scale, | ||
| i8::MIN | ||
| ); | ||
| } | ||
| let precision = if scale > 0 { | ||
| // arrow-rs requires the precision to include the positive scale. | ||
| // See <https://github.com/apache/arrow-rs/blob/123045cc766d42d1eb06ee8bb3f09e39ea995ddc/arrow-array/src/types.rs#L1230> | ||
| std::cmp::max(digits, scale.unsigned_abs()) | ||
| } else { | ||
| digits | ||
| }; | ||
| if precision <= DECIMAL128_MAX_PRECISION as u64 { | ||
| let val = int_val.to_i128().ok_or_else(|| { | ||
| // Failures are unexpected here as we have already checked the precision | ||
| internal_datafusion_err!( | ||
| "Unexpected overflow when converting {} to i128", | ||
| int_val | ||
| ) | ||
| })?; | ||
| Ok(Expr::Literal(ScalarValue::Decimal128( | ||
| Some(val), | ||
| precision as u8, | ||
| scale as i8, | ||
| ))) | ||
| } else if precision <= DECIMAL256_MAX_PRECISION as u64 { | ||
| let val = bigint_to_i256(&int_val).ok_or_else(|| { | ||
| // Failures are unexpected here as we have already checked the precision | ||
| internal_datafusion_err!( | ||
| "Unexpected overflow when converting {} to i256", | ||
| int_val | ||
| ) | ||
| })?; | ||
| Ok(Expr::Literal(ScalarValue::Decimal256( | ||
| Some(val), | ||
| precision as u8, | ||
| scale as i8, | ||
| ))) | ||
| } else { | ||
| not_impl_err!( | ||
findepi marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| "Decimal precision {} exceeds the maximum supported precision: {}", | ||
| precision, | ||
| DECIMAL256_MAX_PRECISION | ||
| ) | ||
| } | ||
| } | ||
|
|
||
| #[cfg(test)] | ||
|
|
@@ -379,4 +424,79 @@ mod tests { | |
| assert_eq!(output, expect); | ||
| } | ||
| } | ||
|
|
||
| #[test] | ||
| fn test_bigint_to_i256() { | ||
| let cases = [ | ||
| (BigInt::from(0), Some(i256::from(0))), | ||
| (BigInt::from(1), Some(i256::from(1))), | ||
| (BigInt::from(-1), Some(i256::from(-1))), | ||
| ( | ||
| BigInt::from_str(i256::MAX.to_string().as_str()).unwrap(), | ||
| Some(i256::MAX), | ||
| ), | ||
| ( | ||
| BigInt::from_str(i256::MIN.to_string().as_str()).unwrap(), | ||
| Some(i256::MIN), | ||
| ), | ||
| ( | ||
| // Can't fit into i256 | ||
| BigInt::from_str((i256::MAX.to_string() + "1").as_str()).unwrap(), | ||
| None, | ||
| ), | ||
| ]; | ||
|
|
||
| for (input, expect) in cases { | ||
| let output = bigint_to_i256(&input); | ||
| assert_eq!(output, expect); | ||
| } | ||
| } | ||
|
|
||
| #[test] | ||
| fn test_parse_decimal() { | ||
| // Supported cases | ||
| let cases = [ | ||
| ("0", ScalarValue::Decimal128(Some(0), 1, 0)), | ||
| ("1", ScalarValue::Decimal128(Some(1), 1, 0)), | ||
| ("123.45", ScalarValue::Decimal128(Some(12345), 5, 2)), | ||
| // Digit count is less than scale | ||
| ("0.001", ScalarValue::Decimal128(Some(1), 3, 3)), | ||
| // Scientific notation | ||
| ("123.456e-2", ScalarValue::Decimal128(Some(123456), 6, 5)), | ||
| // Negative scale | ||
| ("123456e128", ScalarValue::Decimal128(Some(123456), 6, -128)), | ||
| // Decimal256 | ||
| ( | ||
| &("9".repeat(39) + "." + "99999"), | ||
| ScalarValue::Decimal256( | ||
| Some(i256::from_string(&"9".repeat(44)).unwrap()), | ||
| 44, | ||
| 5, | ||
| ), | ||
| ), | ||
| ]; | ||
| for (input, expect) in cases { | ||
| let output = parse_decimal(input, true).unwrap(); | ||
| assert_eq!(output, Expr::Literal(expect.arithmetic_negate().unwrap())); | ||
|
|
||
| let output = parse_decimal(input, false).unwrap(); | ||
| assert_eq!(output, Expr::Literal(expect)); | ||
| } | ||
|
|
||
| // scale < i8::MIN | ||
| assert_eq!( | ||
| parse_decimal("1e129", false) | ||
| .unwrap_err() | ||
| .strip_backtrace(), | ||
| "This feature is not implemented: Decimal scale -129 exceeds the minimum supported scale: -128" | ||
| ); | ||
|
|
||
| // Unsupported precision | ||
| assert_eq!( | ||
| parse_decimal(&"1".repeat(77), false) | ||
| .unwrap_err() | ||
| .strip_backtrace(), | ||
| "This feature is not implemented: Decimal precision 77 exceeds the maximum supported precision: 76" | ||
| ); | ||
| } | ||
| } | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
please add unit tests for this function
it should be easy since both BigInt and i256 can convert to/from str
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Added, thank you @findepi