Skip to content

Commit 94d0f32

Browse files
author
B Vadlamani
committed
improve_benchmark_remove_unwanted_branching
1 parent 5abce5f commit 94d0f32

1 file changed

Lines changed: 180 additions & 25 deletions

File tree

  • native/spark-expr/src/conversion_funcs

‎native/spark-expr/src/conversion_funcs/cast.rs‎

Lines changed: 180 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1967,9 +1967,88 @@ fn cast_string_to_int_with_range_check(
19671967
/// Equivalent to
19681968
/// - org.apache.spark.unsafe.types.UTF8String.toInt(IntWrapper intWrapper, boolean allowDecimal)
19691969
/// - org.apache.spark.unsafe.types.UTF8String.toLong(LongWrapper longWrapper, boolean allowDecimal)
1970-
fn do_cast_string_to_int<T: Integer + CheckedSub + CheckedNeg + From<u8> + Copy>(
1970+
fn do_parse_string_to_int_legacy<T: Integer + CheckedSub + CheckedNeg + From<u8> + Copy>(
1971+
str: &str,
1972+
min_value: T,
1973+
) -> SparkResult<Option<T>> {
1974+
let bytes = str.as_bytes();
1975+
let mut start = 0;
1976+
let mut end = bytes.len();
1977+
1978+
while start < end && bytes[start].is_ascii_whitespace() {
1979+
start += 1;
1980+
}
1981+
while end > start && bytes[end - 1].is_ascii_whitespace() {
1982+
end -= 1;
1983+
}
1984+
1985+
if start == end {
1986+
return Ok(None);
1987+
}
1988+
let trimmed_bytes = &bytes[start..end];
1989+
let len = trimmed_bytes.len();
1990+
let mut result: T = T::zero();
1991+
let mut idx = 0;
1992+
let first_char = trimmed_bytes[0];
1993+
let negative = first_char == b'-';
1994+
if negative || first_char == b'+' {
1995+
idx = 1;
1996+
if len == 1 {
1997+
return Ok(None);
1998+
}
1999+
}
2000+
2001+
let radix = T::from(10_u8);
2002+
let stop_value = min_value / radix;
2003+
let mut parse_sign_and_digits = true;
2004+
2005+
for &ch in &trimmed_bytes[idx..] {
2006+
if parse_sign_and_digits {
2007+
if ch == b'.' {
2008+
// truncate decimal in legacy mode
2009+
parse_sign_and_digits = false;
2010+
continue;
2011+
}
2012+
2013+
if !ch.is_ascii_digit() {
2014+
return Ok(None);
2015+
}
2016+
2017+
let digit: T = T::from(ch - b'0');
2018+
2019+
if result < stop_value {
2020+
return Ok(None);
2021+
}
2022+
let v = result * radix;
2023+
match v.checked_sub(&digit) {
2024+
Some(x) if x <= T::zero() => result = x,
2025+
_ => {
2026+
return Ok(None);
2027+
}
2028+
}
2029+
} else {
2030+
if !ch.is_ascii_digit() {
2031+
return Ok(None);
2032+
}
2033+
}
2034+
}
2035+
2036+
if !negative {
2037+
if let Some(neg) = result.checked_neg() {
2038+
if neg < T::zero() {
2039+
return Ok(None);
2040+
}
2041+
result = neg;
2042+
} else {
2043+
return Ok(None);
2044+
}
2045+
}
2046+
2047+
Ok(Some(result))
2048+
}
2049+
2050+
fn do_parse_string_to_int_ansi<T: Integer + CheckedSub + CheckedNeg + From<u8> + Copy>(
19712051
str: &str,
1972-
eval_mode: EvalMode,
19732052
type_name: &str,
19742053
min_value: T,
19752054
) -> SparkResult<Option<T>> {
@@ -1985,7 +2064,7 @@ fn do_cast_string_to_int<T: Integer + CheckedSub + CheckedNeg + From<u8> + Copy>
19852064
}
19862065

19872066
if start == end {
1988-
return none_or_err(eval_mode, type_name, str);
2067+
return Err(invalid_value(str, "STRING", type_name));
19892068
}
19902069
let trimmed_bytes = &bytes[start..end];
19912070
let len = trimmed_bytes.len();
@@ -1996,7 +2075,7 @@ fn do_cast_string_to_int<T: Integer + CheckedSub + CheckedNeg + From<u8> + Copy>
19962075
if negative || first_char == b'+' {
19972076
idx = 1;
19982077
if len == 1 {
1999-
return none_or_err(eval_mode, type_name, str);
2078+
return Err(invalid_value(str, "STRING", type_name));
20002079
}
20012080
}
20022081

@@ -2007,61 +2086,137 @@ fn do_cast_string_to_int<T: Integer + CheckedSub + CheckedNeg + From<u8> + Copy>
20072086
for &ch in &trimmed_bytes[idx..] {
20082087
if parse_sign_and_digits {
20092088
if ch == b'.' {
2010-
if eval_mode == EvalMode::Legacy {
2011-
// truncate decimal in legacy mode
2012-
parse_sign_and_digits = false;
2013-
continue;
2014-
} else {
2015-
return none_or_err(eval_mode, type_name, str);
2089+
return Err(invalid_value(str, "STRING", type_name));
2090+
}
2091+
2092+
if !ch.is_ascii_digit() {
2093+
return Err(invalid_value(str, "STRING", type_name));
2094+
}
2095+
2096+
let digit: T = T::from(ch - b'0');
2097+
2098+
if result < stop_value {
2099+
return Err(invalid_value(str, "STRING", type_name));
2100+
}
2101+
let v = result * radix;
2102+
match v.checked_sub(&digit) {
2103+
Some(x) if x <= T::zero() => result = x,
2104+
_ => {
2105+
return Err(invalid_value(str, "STRING", type_name));
20162106
}
20172107
}
2108+
} else {
2109+
if !ch.is_ascii_digit() {
2110+
return Err(invalid_value(str, "STRING", type_name));
2111+
}
2112+
}
2113+
}
2114+
2115+
if !negative {
2116+
if let Some(neg) = result.checked_neg() {
2117+
if neg < T::zero() {
2118+
return Err(invalid_value(str, "STRING", type_name));
2119+
}
2120+
result = neg;
2121+
} else {
2122+
return Err(invalid_value(str, "STRING", type_name));
2123+
}
2124+
}
2125+
2126+
Ok(Some(result))
2127+
}
2128+
2129+
fn do_parse_string_to_int_try<T: Integer + CheckedSub + CheckedNeg + From<u8> + Copy>(
2130+
str: &str,
2131+
min_value: T,
2132+
) -> SparkResult<Option<T>> {
2133+
let bytes = str.as_bytes();
2134+
let mut start = 0;
2135+
let mut end = bytes.len();
2136+
2137+
while start < end && bytes[start].is_ascii_whitespace() {
2138+
start += 1;
2139+
}
2140+
while end > start && bytes[end - 1].is_ascii_whitespace() {
2141+
end -= 1;
2142+
}
2143+
2144+
if start == end {
2145+
return Ok(None);
2146+
}
2147+
let trimmed_bytes = &bytes[start..end];
2148+
let len = trimmed_bytes.len();
2149+
let mut result: T = T::zero();
2150+
let mut idx = 0;
2151+
let first_char = trimmed_bytes[0];
2152+
let negative = first_char == b'-';
2153+
if negative || first_char == b'+' {
2154+
idx = 1;
2155+
if len == 1 {
2156+
return Ok(None);
2157+
}
2158+
}
2159+
2160+
let radix = T::from(10_u8);
2161+
let stop_value = min_value / radix;
2162+
let mut parse_sign_and_digits = true;
2163+
2164+
for &ch in &trimmed_bytes[idx..] {
2165+
if parse_sign_and_digits {
2166+
if ch == b'.' {
2167+
return Ok(None);
2168+
}
20182169

20192170
if !ch.is_ascii_digit() {
2020-
return none_or_err(eval_mode, type_name, str);
2171+
return Ok(None);
20212172
}
20222173

2023-
// Direct conversion: u8 digit (0-9) → T
20242174
let digit: T = T::from(ch - b'0');
20252175

2026-
// We are going to process the new digit and accumulate the result. However, before
2027-
// doing this, if the result is already smaller than the
2028-
// stopValue(Integer.MIN_VALUE / radix), then result * 10 will definitely be
2029-
// smaller than minValue, and we can stop
20302176
if result < stop_value {
2031-
return none_or_err(eval_mode, type_name, str);
2177+
return Ok(None);
20322178
}
2033-
// Since the previous result is greater than or equal to stopValue(Integer.MIN_VALUE /
2034-
// radix), we can just use `result > 0` to check overflow. If result
2035-
// overflows, we should stop
20362179
let v = result * radix;
20372180
match v.checked_sub(&digit) {
20382181
Some(x) if x <= T::zero() => result = x,
20392182
_ => {
2040-
return none_or_err(eval_mode, type_name, str);
2183+
return Ok(None);
20412184
}
20422185
}
20432186
} else {
2044-
// make sure fractional digits are valid digits but ignore them
20452187
if !ch.is_ascii_digit() {
2046-
return none_or_err(eval_mode, type_name, str);
2188+
return Ok(None);
20472189
}
20482190
}
20492191
}
20502192

20512193
if !negative {
20522194
if let Some(neg) = result.checked_neg() {
20532195
if neg < T::zero() {
2054-
return none_or_err(eval_mode, type_name, str);
2196+
return Ok(None);
20552197
}
20562198
result = neg;
20572199
} else {
2058-
return none_or_err(eval_mode, type_name, str);
2200+
return Ok(None);
20592201
}
20602202
}
20612203

20622204
Ok(Some(result))
20632205
}
20642206

2207+
fn do_cast_string_to_int<T: Integer + CheckedSub + CheckedNeg + From<u8> + Copy>(
2208+
str: &str,
2209+
eval_mode: EvalMode,
2210+
type_name: &str,
2211+
min_value: T,
2212+
) -> SparkResult<Option<T>> {
2213+
match eval_mode {
2214+
EvalMode::Legacy => do_parse_string_to_int_legacy(str, min_value),
2215+
EvalMode::Ansi => do_parse_string_to_int_ansi(str, type_name, min_value),
2216+
EvalMode::Try => do_parse_string_to_int_try(str, min_value),
2217+
}
2218+
}
2219+
20652220
fn cast_string_to_decimal(
20662221
array: &ArrayRef,
20672222
to_type: &DataType,

0 commit comments

Comments
 (0)