Rich-T-kid commented on code in PR #25874:
URL: https://github.com/apache/datafusion/pull/25874#discussion_r4360301985
##########
datafusion/functions/src/unicode/substr.rs:
##########
@@ -143,13 +164,83 @@ fn substr(args: &[ArrayRef]) -> Result<ArrayRef> {
let string_array = args[0].as_string_view();
string_view_substr(string_array, &args[1..])
}
+ DataType::Dictionary(_, _) => {
+ let array = &args[0];
+ downcast_dictionary_array! {
+ array => dictionary_substr(array, &args[1..]),
+ _ => unreachable!("substr matched a dictionary array"),
+ }
+ }
other => exec_err!(
"Unsupported data type {other:?} for function substr,\
expected Utf8View, Utf8 or LargeUtf8."
),
}
}
+fn dictionary_substr<K: ArrowDictionaryKeyType>(
+ dictionary: &DictionaryArray<K>,
+ args: &[ArrayRef],
+) -> Result<ArrayRef> {
+ let start_array = as_int64_array(&args[0])?;
+ let count_array = args.get(1).map(|a| as_int64_array(a)).transpose()?;
+
+ let mut combinations = HashMap::new();
+ let mut value_indices = Vec::new();
+ let mut starts = Vec::new();
+ let mut counts = count_array.map(|_| Vec::new());
+ let mut keys = PrimitiveBuilder::<K>::with_capacity(dictionary.len());
+
+ for i in 0..dictionary.len() {
+ let Some(value_index) = dictionary.key(i) else {
+ keys.append_null();
+ continue;
+ };
+ if dictionary.values().is_null(value_index)
+ || start_array.is_null(i)
+ || count_array.is_some_and(|a| a.is_null(i))
+ {
+ keys.append_null();
+ continue;
+ }
+
+ let start = start_array.value(i);
+ let count = count_array.map(|a| a.value(i));
+ let combination = (value_index, start, count);
+ let next_key = combinations.len();
+ let key = *combinations.entry(combination).or_insert_with(|| {
+ value_indices.push(value_index as u64);
+ starts.push(start);
+ if let (Some(counts), Some(count)) = (&mut counts, count) {
+ counts.push(count);
+ }
+ next_key
+ });
+ let key = K::Native::from_usize(key).ok_or_else(|| {
Review Comment:
updated the PR to address 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: [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]
