comphead commented on code in PR #2099:
URL: https://github.com/apache/datafusion-comet/pull/2099#discussion_r2311130484
##########
native/spark-expr/src/static_invoke/char_varchar_utils/read_side_padding.rs:
##########
@@ -73,42 +124,71 @@ fn spark_read_side_padding2(
fn spark_read_side_padding_internal<T: OffsetSizeTrait>(
array: &ArrayRef,
- length: i32,
truncate: bool,
+ pad_type: ColumnarValue,
) -> Result<ColumnarValue, DataFusionError> {
let string_array = as_generic_string_array::<T>(array)?;
- let length = 0.max(length) as usize;
- let space_string = " ".repeat(length);
+ match pad_type {
+ ColumnarValue::Array(array_int) => {
+ let int_pad_array = array_int.as_primitive::<Int32Type>();
- let mut builder =
- GenericStringBuilder::<T>::with_capacity(string_array.len(),
string_array.len() * length);
+ let mut builder = GenericStringBuilder::<T>::with_capacity(
+ string_array.len(),
+ string_array.len() * int_pad_array.len(),
+ );
- for string in string_array.iter() {
- match string {
- Some(string) => {
- // It looks Spark's UTF8String is closer to chars rather than
graphemes
- // https://stackoverflow.com/a/46290728
- let char_len = string.chars().count();
- if length <= char_len {
- if truncate {
- let idx = string
- .char_indices()
- .nth(length)
- .map(|(i, _)| i)
- .unwrap_or(string.len());
- builder.append_value(&string[..idx]);
- } else {
- builder.append_value(string);
- }
- } else {
- // write_str updates only the value buffer, not null nor
offset buffer
- // This is convenient for concatenating str(s)
- builder.write_str(string)?;
- builder.append_value(&space_string[char_len..]);
+ for (string, length) in string_array.iter().zip(int_pad_array) {
+ match string {
+ Some(string) => builder.append_value(add_padding_string(
+ string.parse().unwrap(),
+ length.unwrap() as usize,
+ truncate,
+ )),
+ _ => builder.append_null(),
+ }
+ }
+ Ok(ColumnarValue::Array(Arc::new(builder.finish())))
+ }
+ ColumnarValue::Scalar(const_pad_length) => {
+ let length = 0.max(i32::try_from(const_pad_length)?) as usize;
+
+ let mut builder = GenericStringBuilder::<T>::with_capacity(
+ string_array.len(),
+ string_array.len() * length,
+ );
+
+ for string in string_array.iter() {
+ match string {
+ Some(string) => builder.append_value(add_padding_string(
+ string.parse().unwrap(),
+ length,
+ truncate,
+ )),
+ _ => builder.append_null(),
}
}
- _ => builder.append_null(),
+ Ok(ColumnarValue::Array(Arc::new(builder.finish())))
+ }
+ }
+}
+
+fn add_padding_string(string: String, length: usize, truncate: bool) -> String
{
Review Comment:
perhaps we can think of impl like
```
fn add_padding_string(input: String, length: usize, truncate: bool) ->
String {
let char_len = input.chars().count();
if char_len >= length {
if truncate {
// Take the first `length` chars safely
input.chars().take(length).collect()
} else {
input
}
} else {
// Pad with only the needed spaces
let padding = " ".repeat(length - char_len);
input + &padding
}
}
```
so we don't allocate spaces if its not needed
no unwrap
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
