jayzhan211 commented on code in PR #25874:
URL: https://github.com/apache/datafusion/pull/25874#discussion_r4355803914
##########
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:
With per-row `start`/`length`, the output dictionary can need one value per
row, but it keeps the input key type. For `Int8`/`UInt8` keys this fails as
soon as a batch has more than 127/255 distinct combinations, while main returns
the result (it decodes first):
```sql
create table t8 as select arrow_cast(arrow_cast('hello world', 'Utf8'),
'Dictionary(Int8, Utf8)') c, value s from generate_series(1, 200);
select count(*), count(distinct substr(c, s)) from t8;
-- main: 200 | 12
-- this PR: Execution error: substr result has too many dictionary values
for key type Int8
```
The return type can't depend on whether the arguments are scalars, so one
fix is to only preserve the dictionary for key types that can index any batch,
and decode otherwise. I checked locally that this fixes the repro and that
`string/` slt still passes. Could you also add the repro above to
`dictionary_utf8.slt`?
```diff
fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
- Ok(arg_types[0].clone())
+ Ok(match &arg_types[0] {
+ DataType::Dictionary(key_type, value_type) if
!is_wide_key(key_type) => {
+ value_type.as_ref().clone()
+ }
+ other => other.clone(),
+ })
}
fn invoke_with_args(&self, args: ScalarFunctionArgs) ->
Result<ColumnarValue> {
+ if let DataType::Dictionary(key_type, value_type) =
args.args[0].data_type()
+ && !is_wide_key(&key_type)
+ {
+ // Per-row start/length can produce more distinct values than a
+ // narrow key type can index, so decode instead.
+ let mut args = args.args;
+ args[0] = args[0].cast_to(&value_type, None)?;
+ return make_scalar_function(substr, vec![])(&args);
+ }
+
```
```rs
/// Whether `key_type` can index one dictionary value per row of any batch.
fn is_wide_key(key_type: &DataType) -> bool {
matches!(
key_type,
DataType::Int32 | DataType::UInt32 | DataType::Int64 |
DataType::UInt64
)
}
```
--
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]
