andygrove commented on code in PR #307: URL: https://github.com/apache/datafusion-comet/pull/307#discussion_r1579561300
########## core/src/execution/datafusion/expressions/cast.rs: ########## @@ -142,6 +226,281 @@ impl Cast { } } +fn cast_string_to_i8(str: &str, eval_mode: EvalMode) -> CometResult<Option<i8>> { + Ok(cast_string_to_int_with_range_check( + str, + eval_mode, + "TINYINT", + i8::MIN as i32, + i8::MAX as i32, + )? + .map(|v| v as i8)) +} + +fn cast_string_to_i16(str: &str, eval_mode: EvalMode) -> CometResult<Option<i16>> { + Ok(cast_string_to_int_with_range_check( + str, + eval_mode, + "SMALLINT", + i16::MIN as i32, + i16::MAX as i32, + )? + .map(|v| v as i16)) +} + +fn cast_string_to_i32(str: &str, eval_mode: EvalMode) -> CometResult<Option<i32>> { + let mut accum = CastStringToInt32::default(); + do_cast_string_to_int(&mut accum, str, eval_mode, "INT")?; + Ok(accum.result) +} + +fn cast_string_to_i64(str: &str, eval_mode: EvalMode) -> CometResult<Option<i64>> { + let mut accum = CastStringToInt64::default(); + do_cast_string_to_int(&mut accum, str, eval_mode, "BIGINT")?; + Ok(accum.result) +} + +fn cast_string_to_int_with_range_check( + str: &str, + eval_mode: EvalMode, + type_name: &str, + min: i32, + max: i32, +) -> CometResult<Option<i32>> { + let mut accum = CastStringToInt32::default(); + do_cast_string_to_int(&mut accum, str, eval_mode, type_name)?; + match accum.result { + None => Ok(None), + Some(v) if v >= min && v <= max => Ok(Some(v)), + _ if eval_mode == EvalMode::Ansi => Err(invalid_value(str, "STRING", type_name)), + _ => Ok(None), + } +} + +/// We support parsing strings to i32 and i64 to match Spark's logic. Support for i8 and i16 is +/// implemented by first parsing as i32 and then downcasting. The CastStringToInt trait is +/// introduced so that we can have the parsing logic delegate either to an i32 or i64 accumulator +/// and avoid the need to use macros here. +trait CastStringToInt { + fn accumulate( + &mut self, + eval_mode: EvalMode, + type_name: &str, + str: &str, + digit: u32, + ) -> CometResult<()>; + + fn reset(&mut self); + + fn finish( + &mut self, + eval_mode: EvalMode, + type_name: &str, + str: &str, + negative: bool, + ) -> CometResult<()>; +} +struct CastStringToInt32 { + negative: bool, + result: Option<i32>, + radix: i32, +} + +impl Default for CastStringToInt32 { + fn default() -> Self { + Self { + negative: false, + result: Some(0), + radix: 10, + } + } +} + +impl CastStringToInt for CastStringToInt32 { + fn accumulate( + &mut self, + eval_mode: EvalMode, + type_name: &str, + str: &str, + digit: u32, + ) -> CometResult<()> { + if self.result.is_some() && self.result.unwrap() < i32::MIN / self.radix { + self.reset(); + return none_or_err(eval_mode, type_name, str); + } + self.result = Some(self.result.unwrap_or(0) * self.radix - digit as i32); + if self.result.unwrap() > 0 { + self.reset(); + return none_or_err(eval_mode, type_name, str); + } + Ok(()) + } + fn reset(&mut self) { + self.result = None; + } + + fn finish( + &mut self, + eval_mode: EvalMode, + type_name: &str, + str: &str, + negative: bool, + ) -> CometResult<()> { + if self.result.is_some() && !negative { + self.result = Some(-self.result.unwrap()); + if self.result.unwrap() < 0 { + return none_or_err(eval_mode, type_name, str); + } + } + Ok(()) + } +} + +struct CastStringToInt64 { + negative: bool, + result: Option<i64>, + radix: i64, +} + +impl Default for CastStringToInt64 { + fn default() -> Self { + Self { + negative: false, + result: Some(0), + radix: 10, + } + } +} + +impl CastStringToInt for CastStringToInt64 { + fn accumulate( + &mut self, + eval_mode: EvalMode, + type_name: &str, + str: &str, + digit: u32, + ) -> CometResult<()> { + if self.result.unwrap_or(0) < i64::MIN / self.radix { + self.reset(); + return none_or_err(eval_mode, type_name, str); + } + self.result = Some(self.result.unwrap_or(0) * self.radix - digit as i64); + if self.result.unwrap() > 0 { + self.reset(); + return none_or_err(eval_mode, type_name, str); + } + Ok(()) + } + + fn reset(&mut self) { + self.result = None; + } + + fn finish( + &mut self, + eval_mode: EvalMode, + type_name: &str, + str: &str, + negative: bool, + ) -> CometResult<()> { + if self.result.is_some() && !negative { + self.result = Some(-self.result.unwrap()); + if self.result.unwrap() < 0 { + return none_or_err(eval_mode, type_name, str); + } + } + Ok(()) + } +} + +fn do_cast_string_to_int( + accumulator: &mut dyn CastStringToInt, + str: &str, + eval_mode: EvalMode, + type_name: &str, +) -> CometResult<()> { + let chars: Vec<char> = str.chars().collect(); + let mut i = 0; + let mut end = chars.len(); + + // skip leading whitespace + while i < end && chars[i].is_whitespace() { + i += 1; + } + + // skip trailing whitespace + while end > i && chars[end - 1].is_whitespace() { + end -= 1; + } + + // check for empty string + if i == end { + accumulator.reset(); + return Ok(()); + } + + // skip + or - + let negative = chars[0] == '-'; + if negative || chars[0] == '+' { Review Comment: Thank you! The code was originally trimming the string before this point and I missed updating this when I removed the trim. I have now fixed this. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org For additional commands, e-mail: github-h...@datafusion.apache.org