You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "rluvaton (via GitHub)" <gi...@apache.org> on 2024/12/17 15:09:22 UTC
[PR] feat(arrow-select): make list of dictionary merge dictionary keys [arrow-rs]
rluvaton opened a new pull request, #6893:
URL: https://github.com/apache/arrow-rs/pull/6893
# Which issue does this PR close?
Closes #6888
# What changes are included in this PR?
\-
# Are there any user-facing changes?
No
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1932194273
##########
arrow-select/src/concat.rs:
##########
@@ -129,6 +129,54 @@ fn concat_dictionaries<K: ArrowDictionaryKeyType>(
Ok(Arc::new(array))
}
+fn concat_lists<OffsetSize: OffsetSizeTrait>(
+ arrays: &[&dyn Array],
+ field: &FieldRef,
+) -> Result<ArrayRef, ArrowError> {
+ let mut output_len = 0;
+ let mut list_has_nulls = false;
+
+ let lists = arrays
+ .iter()
+ .map(|x| x.as_list::<OffsetSize>())
+ .inspect(|l| {
+ output_len += l.len();
+ list_has_nulls |= l.null_count() != 0;
+ })
+ .collect::<Vec<_>>();
+
+ let lists_nulls = list_has_nulls.then(|| {
+ let mut nulls = BooleanBufferBuilder::new(output_len);
+ for l in &lists {
+ match l.nulls() {
+ Some(n) => nulls.append_buffer(n.inner()),
+ None => nulls.append_n(l.len(), true),
+ }
+ }
+ NullBuffer::new(nulls.finish())
+ });
+
+ let values: Vec<&dyn Array> = lists
+ .iter()
+ .map(|x| x.values().as_ref())
Review Comment:
I think the bug is here, it needs to slice values() by the first offset
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1889941209
##########
arrow-buffer/src/buffer/offset.rs:
##########
@@ -162,6 +163,44 @@ impl<O: ArrowNativeType> OffsetBuffer<O> {
}
}
+impl<O: ArrowNativeType + Integer + Copy> OffsetBuffer<O> {
Review Comment:
The conversion from vec is zero copy
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1889289094
##########
arrow-select/src/concat.rs:
##########
@@ -169,6 +264,21 @@ pub fn concat(arrays: &[&dyn Array]) -> Result<ArrayRef, ArrowError> {
_ => unreachable!("illegal dictionary key type {k}")
};
} else {
+ if let DataType::List(field) = d {
+ if let DataType::Dictionary(k, _) = field.data_type() {
+ downcast_integer! {
+ k.as_ref() => (list_dict_helper, i32, arrays),
+ _ => unreachable!("illegal dictionary key type {k}")
+ };
+ }
+ } else if let DataType::LargeList(field) = d {
+ if let DataType::Dictionary(k, _) = field.data_type() {
+ downcast_integer! {
+ k.as_ref() => (list_dict_helper, i64, arrays),
+ _ => unreachable!("illegal dictionary key type {k}")
+ };
+ }
+ }
Review Comment:
This is quite a lot of codegen for what is a relatively niche use-case, I'm not sure we want to specialize 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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): make list of dictionary merge dictionary keys [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1888681995
##########
arrow-select/src/concat.rs:
##########
@@ -851,4 +1018,169 @@ mod tests {
assert_eq!(array.null_count(), 10);
assert_eq!(array.logical_null_count(), 10);
}
+
+ #[test]
+ fn concat_dictionary_list_array_simple() {
+ let scalars = vec![
+ create_single_row_list_of_dict(vec![Some("a")]),
+ create_single_row_list_of_dict(vec![Some("a")]),
+ create_single_row_list_of_dict(vec![Some("b")]),
+ ];
+
+ let arrays = scalars.iter().map(|a| a as &(dyn Array)).collect::<Vec<_>>();
+ let concat_res = concat(arrays.as_slice()).unwrap();
+
+ let expected_list = create_list_of_dict(vec![
+ // Row 1
+ Some(vec![Some("a")]),
+ Some(vec![Some("a")]),
+ Some(vec![Some("b")]),
+ ]);
+
+ let list = concat_res.as_list::<i32>();
+
+ // Assert that the list is equal to the expected list
+ list.iter().zip(expected_list.iter()).for_each(|(a, b)| {
+ assert_eq!(a, b);
+ });
+
+ let dict = list
+ .values()
+ .as_dictionary::<Int32Type>()
+ .downcast_dict::<StringArray>()
+ .unwrap();
+ println!("{:?}", dict);
+
+ assert_dictionary_has_unique_values::<_, StringArray>(
+ list.values().as_dictionary::<Int32Type>(),
+ );
+ }
+
+ #[test]
+ fn concat_dictionary_list_array_with_multiple_rows() {
Review Comment:
this test fails because there are more keys than there are values, which is weird, see why here:
- https://github.com/apache/arrow-rs/issues/6888#issuecomment-2548648179
Should I remove this test or what?
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1894654161
##########
arrow-select/src/concat.rs:
##########
@@ -129,12 +130,106 @@ fn concat_dictionaries<K: ArrowDictionaryKeyType>(
Ok(Arc::new(array))
}
+fn concat_list_of_dictionaries<OffsetSize: OffsetSizeTrait, K: ArrowDictionaryKeyType>(
Review Comment:
updated, please let me know it that was what you meant, I think it will be slower this way, no?
##########
arrow-select/src/concat.rs:
##########
@@ -851,4 +1018,169 @@ mod tests {
assert_eq!(array.null_count(), 10);
assert_eq!(array.logical_null_count(), 10);
}
+
+ #[test]
+ fn concat_dictionary_list_array_simple() {
+ let scalars = vec![
+ create_single_row_list_of_dict(vec![Some("a")]),
+ create_single_row_list_of_dict(vec![Some("a")]),
+ create_single_row_list_of_dict(vec![Some("b")]),
+ ];
+
+ let arrays = scalars.iter().map(|a| a as &(dyn Array)).collect::<Vec<_>>();
+ let concat_res = concat(arrays.as_slice()).unwrap();
+
+ let expected_list = create_list_of_dict(vec![
+ // Row 1
+ Some(vec![Some("a")]),
+ Some(vec![Some("a")]),
+ Some(vec![Some("b")]),
+ ]);
+
+ let list = concat_res.as_list::<i32>();
+
+ // Assert that the list is equal to the expected list
+ list.iter().zip(expected_list.iter()).for_each(|(a, b)| {
+ assert_eq!(a, b);
+ });
+
+ let dict = list
+ .values()
+ .as_dictionary::<Int32Type>()
+ .downcast_dict::<StringArray>()
+ .unwrap();
+ println!("{:?}", dict);
+
+ assert_dictionary_has_unique_values::<_, StringArray>(
+ list.values().as_dictionary::<Int32Type>(),
+ );
+ }
+
+ #[test]
+ fn concat_dictionary_list_array_with_multiple_rows() {
+ let scalars = vec![
+ create_list_of_dict(vec![
+ // Row 1
+ Some(vec![Some("a"), Some("c")]),
+ // Row 2
+ None,
+ // Row 3
+ Some(vec![Some("f"), Some("g"), None]),
+ // Row 4
+ Some(vec![Some("c"), Some("f")]),
+ ]),
+ create_list_of_dict(vec![
+ // Row 1
+ Some(vec![Some("a")]),
+ // Row 2
+ Some(vec![]),
+ // Row 3
+ Some(vec![None, Some("b")]),
+ // Row 4
+ Some(vec![Some("d"), Some("e")]),
+ ]),
+ create_list_of_dict(vec![
+ // Row 1
+ Some(vec![Some("g")]),
+ // Row 2
+ Some(vec![Some("h"), Some("i")]),
+ // Row 3
+ Some(vec![Some("j"), Some("a")]),
+ // Row 4
+ Some(vec![Some("d"), Some("e")]),
+ ]),
+ ];
+ let arrays = scalars
+ .iter()
+ .map(|a| a as &(dyn Array))
+ .collect::<Vec<_>>();
+ let concat_res = concat(arrays.as_slice()).unwrap();
+
+ let expected_list = create_list_of_dict(vec![
+ // First list:
+
+ // Row 1
+ Some(vec![Some("a"), Some("c")]),
+ // Row 2
+ None,
+ // Row 3
+ Some(vec![Some("f"), Some("g"), None]),
+ // Row 4
+ Some(vec![Some("c"), Some("f")]),
+ // Second list:
+ // Row 1
+ Some(vec![Some("a")]),
+ // Row 2
+ Some(vec![]),
+ // Row 3
+ Some(vec![None, Some("b")]),
+ // Row 4
+ Some(vec![Some("d"), Some("e")]),
+ // Third list:
+
+ // Row 1
+ Some(vec![Some("g")]),
+ // Row 2
+ Some(vec![Some("h"), Some("i")]),
+ // Row 3
+ Some(vec![Some("j"), Some("a")]),
+ // Row 4
+ Some(vec![Some("d"), Some("e")]),
+ ]);
+
+ let list = concat_res.as_list::<i32>();
+
+ // Assert that the list is equal to the expected list
+ list.iter().zip(expected_list.iter()).for_each(|(a, b)| {
+ assert_eq!(a, b);
+ });
+
+ // Assert that the
+ assert_dictionary_has_unique_values::<_, StringArray>(
Review Comment:
removed this test
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1894663564
##########
arrow-select/src/concat.rs:
##########
@@ -129,12 +131,161 @@ fn concat_dictionaries<K: ArrowDictionaryKeyType>(
Ok(Arc::new(array))
}
+fn concat_list_of_dictionaries<OffsetSize: OffsetSizeTrait, K: ArrowDictionaryKeyType>(
+ arrays: &[&dyn Array],
+) -> Result<ArrayRef, ArrowError> {
+ let mut output_len = 0;
+ let mut list_has_nulls = false;
+
+ let lists = arrays
+ .iter()
+ .map(|x| x.as_list::<OffsetSize>())
+ .inspect(|l| {
+ output_len += l.len();
+ list_has_nulls |= l.null_count() != 0;
+ })
+ .collect::<Vec<_>>();
+
+ let mut dictionary_output_len = 0;
+ let dictionaries: Vec<_> = lists
+ .iter()
+ .map(|x| x.values().as_ref().as_dictionary::<K>())
+ .inspect(|d| dictionary_output_len += d.len())
+ .collect();
+
+ if !should_merge_dictionary_values::<K>(&dictionaries, dictionary_output_len) {
+ return concat_fallback(arrays, Capacities::Array(output_len));
+ }
+
+ let merged = merge_dictionary_values(&dictionaries, None)?;
+
+ let lists_nulls = list_has_nulls.then(|| {
+ let mut nulls = BooleanBufferBuilder::new(output_len);
+ for l in &lists {
+ match l.nulls() {
+ Some(n) => nulls.append_buffer(n.inner()),
+ None => nulls.append_n(l.len(), true),
+ }
+ }
+ NullBuffer::new(nulls.finish())
+ });
+
+ // Recompute keys
+ let mut key_values = Vec::with_capacity(dictionary_output_len);
+
+ let mut dictionary_has_nulls = false;
+ for (d, mapping) in dictionaries.iter().zip(merged.key_mappings) {
+ dictionary_has_nulls |= d.null_count() != 0;
+ for key in d.keys().values() {
+ // Use get to safely handle nulls
+ key_values.push(mapping.get(key.as_usize()).copied().unwrap_or_default())
+ }
+ }
+
+ let dictionary_nulls = dictionary_has_nulls.then(|| {
+ let mut nulls = BooleanBufferBuilder::new(dictionary_output_len);
+ for d in &dictionaries {
+ match d.nulls() {
+ Some(n) => nulls.append_buffer(n.inner()),
+ None => nulls.append_n(d.len(), true),
+ }
+ }
+ NullBuffer::new(nulls.finish())
+ });
+
+ let keys = PrimitiveArray::<K>::new(key_values.into(), dictionary_nulls);
+ // Sanity check
+ assert_eq!(keys.len(), dictionary_output_len);
+
+ let array = unsafe { DictionaryArray::new_unchecked(keys, merged.values) };
+
+ // Merge value offsets from the lists
+ let all_value_offsets_iterator = lists
+ .iter()
+ .map(|x| x.offsets());
+
+ let value_offset_buffer = merge_value_offsets(all_value_offsets_iterator);
+
+ let builder = ArrayDataBuilder::new(arrays[0].data_type().clone())
+ .len(output_len)
+ .nulls(lists_nulls)
+ // `GenericListArray` must only have 1 buffer
+ .buffers(vec![value_offset_buffer])
+ // `GenericListArray` must only have 1 child_data
+ .child_data(vec![array.to_data()]);
+
+ // TODO - maybe use build_unchecked?
+ let array_data = builder.build()?;
Review Comment:
I'd recommend using GenericListArray::new, ArrayData is a legacy API that is best avoided other than some very specific use-cases
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1889984752
##########
arrow-buffer/src/buffer/offset.rs:
##########
@@ -162,6 +163,44 @@ impl<O: ArrowNativeType> OffsetBuffer<O> {
}
}
+impl<O: ArrowNativeType + Integer + Copy> OffsetBuffer<O> {
Review Comment:
I also tried to make it more vectorized friendly (although not tested and I only have an idea how to do that as my low level knowledge is not great - do the same thing in a tight loop) I knew there was from length but because it can be a hot path I tried to make it as optimized as possible
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1888682550
##########
arrow-select/src/concat.rs:
##########
@@ -129,12 +131,161 @@ fn concat_dictionaries<K: ArrowDictionaryKeyType>(
Ok(Arc::new(array))
}
+fn concat_list_of_dictionaries<OffsetSize: OffsetSizeTrait, K: ArrowDictionaryKeyType>(
+ arrays: &[&dyn Array],
+) -> Result<ArrayRef, ArrowError> {
+ let mut output_len = 0;
+ let mut list_has_nulls = false;
+
+ let lists = arrays
+ .iter()
+ .map(|x| x.as_list::<OffsetSize>())
+ .inspect(|l| {
+ output_len += l.len();
+ list_has_nulls |= l.null_count() != 0;
+ })
+ .collect::<Vec<_>>();
+
+ let mut dictionary_output_len = 0;
+ let dictionaries: Vec<_> = lists
+ .iter()
+ .map(|x| x.values().as_ref().as_dictionary::<K>())
+ .inspect(|d| dictionary_output_len += d.len())
+ .collect();
+
+ if !should_merge_dictionary_values::<K>(&dictionaries, dictionary_output_len) {
+ return concat_fallback(arrays, Capacities::Array(output_len));
+ }
+
+ let merged = merge_dictionary_values(&dictionaries, None)?;
+
+ let lists_nulls = list_has_nulls.then(|| {
+ let mut nulls = BooleanBufferBuilder::new(output_len);
+ for l in &lists {
+ match l.nulls() {
+ Some(n) => nulls.append_buffer(n.inner()),
+ None => nulls.append_n(l.len(), true),
+ }
+ }
+ NullBuffer::new(nulls.finish())
+ });
+
+ // Recompute keys
+ let mut key_values = Vec::with_capacity(dictionary_output_len);
+
+ let mut dictionary_has_nulls = false;
+ for (d, mapping) in dictionaries.iter().zip(merged.key_mappings) {
+ dictionary_has_nulls |= d.null_count() != 0;
+ for key in d.keys().values() {
+ // Use get to safely handle nulls
+ key_values.push(mapping.get(key.as_usize()).copied().unwrap_or_default())
+ }
+ }
+
+ let dictionary_nulls = dictionary_has_nulls.then(|| {
+ let mut nulls = BooleanBufferBuilder::new(dictionary_output_len);
+ for d in &dictionaries {
+ match d.nulls() {
+ Some(n) => nulls.append_buffer(n.inner()),
+ None => nulls.append_n(d.len(), true),
+ }
+ }
+ NullBuffer::new(nulls.finish())
+ });
+
+ let keys = PrimitiveArray::<K>::new(key_values.into(), dictionary_nulls);
+ // Sanity check
+ assert_eq!(keys.len(), dictionary_output_len);
+
+ let array = unsafe { DictionaryArray::new_unchecked(keys, merged.values) };
+
+ // Merge value offsets from the lists
+ let all_value_offsets_iterator = lists
+ .iter()
+ .map(|x| x.offsets());
+
+ let value_offset_buffer = merge_value_offsets(all_value_offsets_iterator);
+
+ let builder = ArrayDataBuilder::new(arrays[0].data_type().clone())
+ .len(output_len)
+ .nulls(lists_nulls)
+ // `GenericListArray` must only have 1 buffer
+ .buffers(vec![value_offset_buffer])
+ // `GenericListArray` must only have 1 child_data
+ .child_data(vec![array.to_data()]);
+
+ // TODO - maybe use build_unchecked?
+ let array_data = builder.build()?;
+
+ let array = GenericListArray::<OffsetSize>::from(array_data);
+ Ok(Arc::new(array))
+}
+
+/// Merge value offsets
+///
+///
+/// if we have the following
+/// [[0, 3, 5], [0, 2, 2, 8], [], [0, 0, 1]]
+/// The output should be
+/// [ 0, 3, 5, 7, 7, 13, 13, 14]
+fn merge_value_offsets<'a, OffsetSize: OffsetSizeTrait, I: Iterator<Item = &'a OffsetBuffer<OffsetSize>>>(offset_buffers_iterator: I) -> Buffer {
Review Comment:
should I move this to be in the offset buffer / scalar buffer / buffer file
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1889300411
##########
arrow-select/src/concat.rs:
##########
@@ -851,4 +1018,169 @@ mod tests {
assert_eq!(array.null_count(), 10);
assert_eq!(array.logical_null_count(), 10);
}
+
+ #[test]
+ fn concat_dictionary_list_array_simple() {
+ let scalars = vec![
+ create_single_row_list_of_dict(vec![Some("a")]),
+ create_single_row_list_of_dict(vec![Some("a")]),
+ create_single_row_list_of_dict(vec![Some("b")]),
+ ];
+
+ let arrays = scalars.iter().map(|a| a as &(dyn Array)).collect::<Vec<_>>();
+ let concat_res = concat(arrays.as_slice()).unwrap();
+
+ let expected_list = create_list_of_dict(vec![
+ // Row 1
+ Some(vec![Some("a")]),
+ Some(vec![Some("a")]),
+ Some(vec![Some("b")]),
+ ]);
+
+ let list = concat_res.as_list::<i32>();
+
+ // Assert that the list is equal to the expected list
+ list.iter().zip(expected_list.iter()).for_each(|(a, b)| {
+ assert_eq!(a, b);
+ });
+
+ let dict = list
+ .values()
+ .as_dictionary::<Int32Type>()
+ .downcast_dict::<StringArray>()
+ .unwrap();
+ println!("{:?}", dict);
+
+ assert_dictionary_has_unique_values::<_, StringArray>(
+ list.values().as_dictionary::<Int32Type>(),
+ );
+ }
+
+ #[test]
+ fn concat_dictionary_list_array_with_multiple_rows() {
+ let scalars = vec![
+ create_list_of_dict(vec![
+ // Row 1
+ Some(vec![Some("a"), Some("c")]),
+ // Row 2
+ None,
+ // Row 3
+ Some(vec![Some("f"), Some("g"), None]),
+ // Row 4
+ Some(vec![Some("c"), Some("f")]),
+ ]),
+ create_list_of_dict(vec![
+ // Row 1
+ Some(vec![Some("a")]),
+ // Row 2
+ Some(vec![]),
+ // Row 3
+ Some(vec![None, Some("b")]),
+ // Row 4
+ Some(vec![Some("d"), Some("e")]),
+ ]),
+ create_list_of_dict(vec![
+ // Row 1
+ Some(vec![Some("g")]),
+ // Row 2
+ Some(vec![Some("h"), Some("i")]),
+ // Row 3
+ Some(vec![Some("j"), Some("a")]),
+ // Row 4
+ Some(vec![Some("d"), Some("e")]),
+ ]),
+ ];
+ let arrays = scalars
+ .iter()
+ .map(|a| a as &(dyn Array))
+ .collect::<Vec<_>>();
+ let concat_res = concat(arrays.as_slice()).unwrap();
+
+ let expected_list = create_list_of_dict(vec![
+ // First list:
+
+ // Row 1
+ Some(vec![Some("a"), Some("c")]),
+ // Row 2
+ None,
+ // Row 3
+ Some(vec![Some("f"), Some("g"), None]),
+ // Row 4
+ Some(vec![Some("c"), Some("f")]),
+ // Second list:
+ // Row 1
+ Some(vec![Some("a")]),
+ // Row 2
+ Some(vec![]),
+ // Row 3
+ Some(vec![None, Some("b")]),
+ // Row 4
+ Some(vec![Some("d"), Some("e")]),
+ // Third list:
+
+ // Row 1
+ Some(vec![Some("g")]),
+ // Row 2
+ Some(vec![Some("h"), Some("i")]),
+ // Row 3
+ Some(vec![Some("j"), Some("a")]),
+ // Row 4
+ Some(vec![Some("d"), Some("e")]),
+ ]);
+
+ let list = concat_res.as_list::<i32>();
+
+ // Assert that the list is equal to the expected list
+ list.iter().zip(expected_list.iter()).for_each(|(a, b)| {
+ assert_eq!(a, b);
+ });
+
+ // Assert that the
+ assert_dictionary_has_unique_values::<_, StringArray>(
Review Comment:
The dictionary recomputation is a best-effort optimisation, it is not guaranteed to eliminate all duplicates (or even run at all)
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1888719695
##########
arrow-select/src/concat.rs:
##########
@@ -129,12 +130,164 @@ fn concat_dictionaries<K: ArrowDictionaryKeyType>(
Ok(Arc::new(array))
}
+fn concat_list_of_dictionaries<OffsetSize: OffsetSizeTrait, K: ArrowDictionaryKeyType>(
+ arrays: &[&dyn Array],
+) -> Result<ArrayRef, ArrowError> {
+ let mut output_len = 0;
+ let mut list_has_nulls = false;
+
+ let lists = arrays
+ .iter()
+ .map(|x| x.as_list::<OffsetSize>())
+ .inspect(|l| {
+ output_len += l.len();
+ list_has_nulls |= l.null_count() != 0;
+ })
+ .collect::<Vec<_>>();
+
+ let mut dictionary_output_len = 0;
+ let dictionaries: Vec<_> = lists
+ .iter()
+ .map(|x| x.values().as_ref().as_dictionary::<K>())
+ .inspect(|d| dictionary_output_len += d.len())
+ .collect();
+
+ if !should_merge_dictionary_values::<K>(&dictionaries, dictionary_output_len) {
+ return concat_fallback(arrays, Capacities::Array(output_len));
+ }
+
+ let merged = merge_dictionary_values(&dictionaries, None)?;
+
+ let lists_nulls = list_has_nulls.then(|| {
+ let mut nulls = BooleanBufferBuilder::new(output_len);
+ for l in &lists {
+ match l.nulls() {
+ Some(n) => nulls.append_buffer(n.inner()),
+ None => nulls.append_n(l.len(), true),
+ }
+ }
+ NullBuffer::new(nulls.finish())
+ });
+
+ // Recompute keys
+ let mut key_values = Vec::with_capacity(dictionary_output_len);
+
+ let mut dictionary_has_nulls = false;
+ for (d, mapping) in dictionaries.iter().zip(merged.key_mappings) {
+ dictionary_has_nulls |= d.null_count() != 0;
+ for key in d.keys().values() {
+ // Use get to safely handle nulls
+ key_values.push(mapping.get(key.as_usize()).copied().unwrap_or_default())
+ }
+ }
+
+ let dictionary_nulls = dictionary_has_nulls.then(|| {
+ let mut nulls = BooleanBufferBuilder::new(dictionary_output_len);
+ for d in &dictionaries {
+ match d.nulls() {
+ Some(n) => nulls.append_buffer(n.inner()),
+ None => nulls.append_n(d.len(), true),
+ }
+ }
+ NullBuffer::new(nulls.finish())
+ });
+
+ let keys = PrimitiveArray::<K>::new(key_values.into(), dictionary_nulls);
+ // Sanity check
+ assert_eq!(keys.len(), dictionary_output_len);
+
+ let array = unsafe { DictionaryArray::new_unchecked(keys, merged.values) };
+
+ // Merge value offsets from the lists
+ let all_value_offsets_iterator = lists.iter().map(|x| x.offsets());
+
+ let value_offset_buffer = merge_value_offsets(all_value_offsets_iterator);
+
+ let builder = ArrayDataBuilder::new(arrays[0].data_type().clone())
+ .len(output_len)
+ .nulls(lists_nulls)
+ // `GenericListArray` must only have 1 buffer
+ .buffers(vec![value_offset_buffer])
+ // `GenericListArray` must only have 1 child_data
+ .child_data(vec![array.to_data()]);
+
+ // TODO - maybe use build_unchecked?
+ let array_data = builder.build()?;
+
+ let array = GenericListArray::<OffsetSize>::from(array_data);
+ Ok(Arc::new(array))
+}
+
+/// Merge value offsets
+///
+///
+/// if we have the following
+/// [[0, 3, 5], [0, 2, 2, 8], [], [0, 0, 1]]
+/// The output should be
+/// [ 0, 3, 5, 7, 7, 13, 13, 14]
+fn merge_value_offsets<
+ 'a,
+ OffsetSize: OffsetSizeTrait,
+ I: Iterator<Item = &'a OffsetBuffer<OffsetSize>>,
+>(
+ offset_buffers_iterator: I,
+) -> Buffer {
+ // 1. Filter out empty lists
+ let mut offset_buffers_iterator = offset_buffers_iterator.filter(|x| !x.is_empty());
+
+ // 2. Get first non-empty list as the starting point
+ let starting_buffer = offset_buffers_iterator.next();
+
+ // 3. If we have only empty lists, return an empty buffer
+ if starting_buffer.is_none() {
+ return Buffer::from(&[]);
+ }
+
+ let starting_buffer = starting_buffer.unwrap();
+
+ let mut offsets_iter: Box<dyn Iterator<Item = OffsetSize>> =
+ Box::new(starting_buffer.iter().copied());
+
+ // 4. Get the last value in the starting buffer as the starting point for the next buffer
+ // Safety: We already filtered out empty lists
+ let mut advance_by = *starting_buffer.last().unwrap();
+
+ // 5. Iterate over the remaining buffers
+ for offset_buffer in offset_buffers_iterator {
+ // 6. Get the last value of the current buffer so we can know how much to advance the next buffer
+ // Safety: We already filtered out empty lists
+ let last_value = *offset_buffer.last().unwrap();
+
+ // 7. Advance the offset buffer by the last value in the previous buffer
+ let offset_buffer_iter = offset_buffer
+ .iter()
+ // Skip the first value as it is the initial offset of 0
+ .skip(1)
+ .map(move |&x| x + advance_by);
+
+ // 8. concat the current buffer with the previous buffer
+ // Chaining keeps the iterator have trusting length
+ offsets_iter = Box::new(offsets_iter.chain(offset_buffer_iter));
Review Comment:
I have stack overflow error here on large number of array to merge, fixing now
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1889693253
##########
arrow-buffer/src/buffer/offset.rs:
##########
@@ -162,6 +163,44 @@ impl<O: ArrowNativeType> OffsetBuffer<O> {
}
}
+impl<O: ArrowNativeType + Integer + Copy> OffsetBuffer<O> {
Review Comment:
I specifically did not wanna use the from length to improve performance (although I may be wrong)
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1890024817
##########
arrow-buffer/src/buffer/offset.rs:
##########
@@ -162,6 +163,44 @@ impl<O: ArrowNativeType> OffsetBuffer<O> {
}
}
+impl<O: ArrowNativeType + Integer + Copy> OffsetBuffer<O> {
Review Comment:
Noted, will try to fix later your comments
Which dyn dispatch here? The inner iterator?
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold merged PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): make list of dictionary merge dictionary keys [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1888683404
##########
arrow-select/src/concat.rs:
##########
@@ -129,12 +131,161 @@ fn concat_dictionaries<K: ArrowDictionaryKeyType>(
Ok(Arc::new(array))
}
+fn concat_list_of_dictionaries<OffsetSize: OffsetSizeTrait, K: ArrowDictionaryKeyType>(
+ arrays: &[&dyn Array],
+) -> Result<ArrayRef, ArrowError> {
+ let mut output_len = 0;
+ let mut list_has_nulls = false;
+
+ let lists = arrays
+ .iter()
+ .map(|x| x.as_list::<OffsetSize>())
+ .inspect(|l| {
+ output_len += l.len();
+ list_has_nulls |= l.null_count() != 0;
+ })
+ .collect::<Vec<_>>();
+
+ let mut dictionary_output_len = 0;
+ let dictionaries: Vec<_> = lists
+ .iter()
+ .map(|x| x.values().as_ref().as_dictionary::<K>())
+ .inspect(|d| dictionary_output_len += d.len())
+ .collect();
+
+ if !should_merge_dictionary_values::<K>(&dictionaries, dictionary_output_len) {
+ return concat_fallback(arrays, Capacities::Array(output_len));
+ }
+
+ let merged = merge_dictionary_values(&dictionaries, None)?;
+
+ let lists_nulls = list_has_nulls.then(|| {
+ let mut nulls = BooleanBufferBuilder::new(output_len);
+ for l in &lists {
+ match l.nulls() {
+ Some(n) => nulls.append_buffer(n.inner()),
+ None => nulls.append_n(l.len(), true),
+ }
+ }
+ NullBuffer::new(nulls.finish())
+ });
+
+ // Recompute keys
+ let mut key_values = Vec::with_capacity(dictionary_output_len);
+
+ let mut dictionary_has_nulls = false;
+ for (d, mapping) in dictionaries.iter().zip(merged.key_mappings) {
+ dictionary_has_nulls |= d.null_count() != 0;
+ for key in d.keys().values() {
+ // Use get to safely handle nulls
+ key_values.push(mapping.get(key.as_usize()).copied().unwrap_or_default())
+ }
+ }
+
+ let dictionary_nulls = dictionary_has_nulls.then(|| {
+ let mut nulls = BooleanBufferBuilder::new(dictionary_output_len);
+ for d in &dictionaries {
+ match d.nulls() {
+ Some(n) => nulls.append_buffer(n.inner()),
+ None => nulls.append_n(d.len(), true),
+ }
+ }
+ NullBuffer::new(nulls.finish())
+ });
+
+ let keys = PrimitiveArray::<K>::new(key_values.into(), dictionary_nulls);
+ // Sanity check
+ assert_eq!(keys.len(), dictionary_output_len);
+
+ let array = unsafe { DictionaryArray::new_unchecked(keys, merged.values) };
+
+ // Merge value offsets from the lists
+ let all_value_offsets_iterator = lists
+ .iter()
+ .map(|x| x.offsets());
+
+ let value_offset_buffer = merge_value_offsets(all_value_offsets_iterator);
+
+ let builder = ArrayDataBuilder::new(arrays[0].data_type().clone())
+ .len(output_len)
+ .nulls(lists_nulls)
+ // `GenericListArray` must only have 1 buffer
+ .buffers(vec![value_offset_buffer])
+ // `GenericListArray` must only have 1 child_data
+ .child_data(vec![array.to_data()]);
+
+ // TODO - maybe use build_unchecked?
+ let array_data = builder.build()?;
Review Comment:
Should I use `build` or `build_unchecked`?
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): make list of dictionary merge dictionary keys [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1888682550
##########
arrow-select/src/concat.rs:
##########
@@ -129,12 +131,161 @@ fn concat_dictionaries<K: ArrowDictionaryKeyType>(
Ok(Arc::new(array))
}
+fn concat_list_of_dictionaries<OffsetSize: OffsetSizeTrait, K: ArrowDictionaryKeyType>(
+ arrays: &[&dyn Array],
+) -> Result<ArrayRef, ArrowError> {
+ let mut output_len = 0;
+ let mut list_has_nulls = false;
+
+ let lists = arrays
+ .iter()
+ .map(|x| x.as_list::<OffsetSize>())
+ .inspect(|l| {
+ output_len += l.len();
+ list_has_nulls |= l.null_count() != 0;
+ })
+ .collect::<Vec<_>>();
+
+ let mut dictionary_output_len = 0;
+ let dictionaries: Vec<_> = lists
+ .iter()
+ .map(|x| x.values().as_ref().as_dictionary::<K>())
+ .inspect(|d| dictionary_output_len += d.len())
+ .collect();
+
+ if !should_merge_dictionary_values::<K>(&dictionaries, dictionary_output_len) {
+ return concat_fallback(arrays, Capacities::Array(output_len));
+ }
+
+ let merged = merge_dictionary_values(&dictionaries, None)?;
+
+ let lists_nulls = list_has_nulls.then(|| {
+ let mut nulls = BooleanBufferBuilder::new(output_len);
+ for l in &lists {
+ match l.nulls() {
+ Some(n) => nulls.append_buffer(n.inner()),
+ None => nulls.append_n(l.len(), true),
+ }
+ }
+ NullBuffer::new(nulls.finish())
+ });
+
+ // Recompute keys
+ let mut key_values = Vec::with_capacity(dictionary_output_len);
+
+ let mut dictionary_has_nulls = false;
+ for (d, mapping) in dictionaries.iter().zip(merged.key_mappings) {
+ dictionary_has_nulls |= d.null_count() != 0;
+ for key in d.keys().values() {
+ // Use get to safely handle nulls
+ key_values.push(mapping.get(key.as_usize()).copied().unwrap_or_default())
+ }
+ }
+
+ let dictionary_nulls = dictionary_has_nulls.then(|| {
+ let mut nulls = BooleanBufferBuilder::new(dictionary_output_len);
+ for d in &dictionaries {
+ match d.nulls() {
+ Some(n) => nulls.append_buffer(n.inner()),
+ None => nulls.append_n(d.len(), true),
+ }
+ }
+ NullBuffer::new(nulls.finish())
+ });
+
+ let keys = PrimitiveArray::<K>::new(key_values.into(), dictionary_nulls);
+ // Sanity check
+ assert_eq!(keys.len(), dictionary_output_len);
+
+ let array = unsafe { DictionaryArray::new_unchecked(keys, merged.values) };
+
+ // Merge value offsets from the lists
+ let all_value_offsets_iterator = lists
+ .iter()
+ .map(|x| x.offsets());
+
+ let value_offset_buffer = merge_value_offsets(all_value_offsets_iterator);
+
+ let builder = ArrayDataBuilder::new(arrays[0].data_type().clone())
+ .len(output_len)
+ .nulls(lists_nulls)
+ // `GenericListArray` must only have 1 buffer
+ .buffers(vec![value_offset_buffer])
+ // `GenericListArray` must only have 1 child_data
+ .child_data(vec![array.to_data()]);
+
+ // TODO - maybe use build_unchecked?
+ let array_data = builder.build()?;
+
+ let array = GenericListArray::<OffsetSize>::from(array_data);
+ Ok(Arc::new(array))
+}
+
+/// Merge value offsets
+///
+///
+/// if we have the following
+/// [[0, 3, 5], [0, 2, 2, 8], [], [0, 0, 1]]
+/// The output should be
+/// [ 0, 3, 5, 7, 7, 13, 13, 14]
+fn merge_value_offsets<'a, OffsetSize: OffsetSizeTrait, I: Iterator<Item = &'a OffsetBuffer<OffsetSize>>>(offset_buffers_iterator: I) -> Buffer {
Review Comment:
should I move this to be in the offset buffer / scalar buffer / buffer file
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1889984752
##########
arrow-buffer/src/buffer/offset.rs:
##########
@@ -162,6 +163,44 @@ impl<O: ArrowNativeType> OffsetBuffer<O> {
}
}
+impl<O: ArrowNativeType + Integer + Copy> OffsetBuffer<O> {
Review Comment:
I also tried to make it more vectorized friendly (although not tested and I only have an idea how to do that as my low level knowledge is not great - do the same thing in a tight loop)
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1890021873
##########
arrow-buffer/src/buffer/offset.rs:
##########
@@ -162,6 +163,44 @@ impl<O: ArrowNativeType> OffsetBuffer<O> {
}
}
+impl<O: ArrowNativeType + Integer + Copy> OffsetBuffer<O> {
Review Comment:
I would recommend a focus on getting something correct and working before trying to optimise it, and then when trying to optimise use empirical benchmarks.
FWIW I suspect this code will optimise very poorly as a result of the dynamic dispatch.
TLDR I would recommend using from_lengths (or OffsetBufferBuilder) unless you can empirically show an alternative approach is sufficiently faster to warrant the additional complexity.
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1888695883
##########
arrow-select/src/concat.rs:
##########
@@ -129,12 +130,164 @@ fn concat_dictionaries<K: ArrowDictionaryKeyType>(
Ok(Arc::new(array))
}
+fn concat_list_of_dictionaries<OffsetSize: OffsetSizeTrait, K: ArrowDictionaryKeyType>(
+ arrays: &[&dyn Array],
+) -> Result<ArrayRef, ArrowError> {
+ let mut output_len = 0;
+ let mut list_has_nulls = false;
+
+ let lists = arrays
+ .iter()
+ .map(|x| x.as_list::<OffsetSize>())
+ .inspect(|l| {
+ output_len += l.len();
+ list_has_nulls |= l.null_count() != 0;
+ })
+ .collect::<Vec<_>>();
+
+ let mut dictionary_output_len = 0;
+ let dictionaries: Vec<_> = lists
+ .iter()
+ .map(|x| x.values().as_ref().as_dictionary::<K>())
+ .inspect(|d| dictionary_output_len += d.len())
+ .collect();
+
+ if !should_merge_dictionary_values::<K>(&dictionaries, dictionary_output_len) {
+ return concat_fallback(arrays, Capacities::Array(output_len));
+ }
+
+ let merged = merge_dictionary_values(&dictionaries, None)?;
+
+ let lists_nulls = list_has_nulls.then(|| {
+ let mut nulls = BooleanBufferBuilder::new(output_len);
+ for l in &lists {
+ match l.nulls() {
+ Some(n) => nulls.append_buffer(n.inner()),
+ None => nulls.append_n(l.len(), true),
+ }
+ }
+ NullBuffer::new(nulls.finish())
+ });
+
+ // Recompute keys
+ let mut key_values = Vec::with_capacity(dictionary_output_len);
+
+ let mut dictionary_has_nulls = false;
+ for (d, mapping) in dictionaries.iter().zip(merged.key_mappings) {
+ dictionary_has_nulls |= d.null_count() != 0;
+ for key in d.keys().values() {
+ // Use get to safely handle nulls
+ key_values.push(mapping.get(key.as_usize()).copied().unwrap_or_default())
+ }
+ }
+
+ let dictionary_nulls = dictionary_has_nulls.then(|| {
+ let mut nulls = BooleanBufferBuilder::new(dictionary_output_len);
+ for d in &dictionaries {
+ match d.nulls() {
+ Some(n) => nulls.append_buffer(n.inner()),
+ None => nulls.append_n(d.len(), true),
+ }
+ }
+ NullBuffer::new(nulls.finish())
+ });
+
+ let keys = PrimitiveArray::<K>::new(key_values.into(), dictionary_nulls);
+ // Sanity check
+ assert_eq!(keys.len(), dictionary_output_len);
+
+ let array = unsafe { DictionaryArray::new_unchecked(keys, merged.values) };
+
+ // Merge value offsets from the lists
+ let all_value_offsets_iterator = lists.iter().map(|x| x.offsets());
+
+ let value_offset_buffer = merge_value_offsets(all_value_offsets_iterator);
+
+ let builder = ArrayDataBuilder::new(arrays[0].data_type().clone())
+ .len(output_len)
+ .nulls(lists_nulls)
+ // `GenericListArray` must only have 1 buffer
+ .buffers(vec![value_offset_buffer])
+ // `GenericListArray` must only have 1 child_data
+ .child_data(vec![array.to_data()]);
+
+ // TODO - maybe use build_unchecked?
+ let array_data = builder.build()?;
+
+ let array = GenericListArray::<OffsetSize>::from(array_data);
+ Ok(Arc::new(array))
+}
+
+/// Merge value offsets
+///
+///
+/// if we have the following
+/// [[0, 3, 5], [0, 2, 2, 8], [], [0, 0, 1]]
+/// The output should be
+/// [ 0, 3, 5, 7, 7, 13, 13, 14]
+fn merge_value_offsets<
Review Comment:
should I move this to be in the offset buffer / scalar buffer / buffer file?
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1889285839
##########
arrow-buffer/src/buffer/offset.rs:
##########
@@ -184,6 +223,119 @@ impl<O: ArrowNativeType> From<OffsetBufferBuilder<O>> for OffsetBuffer<O> {
}
}
+struct MergeBuffersIter<'a, Offset: Integer + Copy> {
Review Comment:
This seems exceptionally complicated for what it is doing, see above for how this could be achieved with a lot less complexity
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1889693253
##########
arrow-buffer/src/buffer/offset.rs:
##########
@@ -162,6 +163,44 @@ impl<O: ArrowNativeType> OffsetBuffer<O> {
}
}
+impl<O: ArrowNativeType + Integer + Copy> OffsetBuffer<O> {
Review Comment:
I specifically did not wanna use the from length to avoid allocating intermediate buffer (the vector inside the from length)
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1894665149
##########
arrow-select/src/concat.rs:
##########
@@ -129,12 +131,161 @@ fn concat_dictionaries<K: ArrowDictionaryKeyType>(
Ok(Arc::new(array))
}
+fn concat_list_of_dictionaries<OffsetSize: OffsetSizeTrait, K: ArrowDictionaryKeyType>(
+ arrays: &[&dyn Array],
+) -> Result<ArrayRef, ArrowError> {
+ let mut output_len = 0;
+ let mut list_has_nulls = false;
+
+ let lists = arrays
+ .iter()
+ .map(|x| x.as_list::<OffsetSize>())
+ .inspect(|l| {
+ output_len += l.len();
+ list_has_nulls |= l.null_count() != 0;
+ })
+ .collect::<Vec<_>>();
+
+ let mut dictionary_output_len = 0;
+ let dictionaries: Vec<_> = lists
+ .iter()
+ .map(|x| x.values().as_ref().as_dictionary::<K>())
+ .inspect(|d| dictionary_output_len += d.len())
+ .collect();
+
+ if !should_merge_dictionary_values::<K>(&dictionaries, dictionary_output_len) {
+ return concat_fallback(arrays, Capacities::Array(output_len));
+ }
+
+ let merged = merge_dictionary_values(&dictionaries, None)?;
+
+ let lists_nulls = list_has_nulls.then(|| {
+ let mut nulls = BooleanBufferBuilder::new(output_len);
+ for l in &lists {
+ match l.nulls() {
+ Some(n) => nulls.append_buffer(n.inner()),
+ None => nulls.append_n(l.len(), true),
+ }
+ }
+ NullBuffer::new(nulls.finish())
+ });
+
+ // Recompute keys
+ let mut key_values = Vec::with_capacity(dictionary_output_len);
+
+ let mut dictionary_has_nulls = false;
+ for (d, mapping) in dictionaries.iter().zip(merged.key_mappings) {
+ dictionary_has_nulls |= d.null_count() != 0;
+ for key in d.keys().values() {
+ // Use get to safely handle nulls
+ key_values.push(mapping.get(key.as_usize()).copied().unwrap_or_default())
+ }
+ }
+
+ let dictionary_nulls = dictionary_has_nulls.then(|| {
+ let mut nulls = BooleanBufferBuilder::new(dictionary_output_len);
+ for d in &dictionaries {
+ match d.nulls() {
+ Some(n) => nulls.append_buffer(n.inner()),
+ None => nulls.append_n(d.len(), true),
+ }
+ }
+ NullBuffer::new(nulls.finish())
+ });
+
+ let keys = PrimitiveArray::<K>::new(key_values.into(), dictionary_nulls);
+ // Sanity check
+ assert_eq!(keys.len(), dictionary_output_len);
+
+ let array = unsafe { DictionaryArray::new_unchecked(keys, merged.values) };
+
+ // Merge value offsets from the lists
+ let all_value_offsets_iterator = lists
+ .iter()
+ .map(|x| x.offsets());
+
+ let value_offset_buffer = merge_value_offsets(all_value_offsets_iterator);
+
+ let builder = ArrayDataBuilder::new(arrays[0].data_type().clone())
+ .len(output_len)
+ .nulls(lists_nulls)
+ // `GenericListArray` must only have 1 buffer
+ .buffers(vec![value_offset_buffer])
+ // `GenericListArray` must only have 1 child_data
+ .child_data(vec![array.to_data()]);
+
+ // TODO - maybe use build_unchecked?
+ let array_data = builder.build()?;
Review Comment:
done
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1894654091
##########
arrow-select/src/concat.rs:
##########
@@ -169,6 +264,21 @@ pub fn concat(arrays: &[&dyn Array]) -> Result<ArrayRef, ArrowError> {
_ => unreachable!("illegal dictionary key type {k}")
};
} else {
+ if let DataType::List(field) = d {
Review Comment:
updated
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1889284118
##########
arrow-buffer/src/buffer/offset.rs:
##########
@@ -162,6 +163,44 @@ impl<O: ArrowNativeType> OffsetBuffer<O> {
}
}
+impl<O: ArrowNativeType + Integer + Copy> OffsetBuffer<O> {
Review Comment:
Perhaps instead of this we could just add
```
pub fn lengths(&self) -> impl Iterator<Item=usize> {
self.0.windows(2).map(|x| x[1].as_usize() - x[0].as_usize())
}
```
Then you'd be able to merge using
```
OffsetBuffer::from_lengths(iter.flat_map(|x| x.lengths()))
```
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1889294686
##########
arrow-select/src/concat.rs:
##########
@@ -169,6 +264,21 @@ pub fn concat(arrays: &[&dyn Array]) -> Result<ArrayRef, ArrowError> {
_ => unreachable!("illegal dictionary key type {k}")
};
} else {
+ if let DataType::List(field) = d {
Review Comment:
It would be better to rework this into a match block
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1888695883
##########
arrow-select/src/concat.rs:
##########
@@ -129,12 +130,164 @@ fn concat_dictionaries<K: ArrowDictionaryKeyType>(
Ok(Arc::new(array))
}
+fn concat_list_of_dictionaries<OffsetSize: OffsetSizeTrait, K: ArrowDictionaryKeyType>(
+ arrays: &[&dyn Array],
+) -> Result<ArrayRef, ArrowError> {
+ let mut output_len = 0;
+ let mut list_has_nulls = false;
+
+ let lists = arrays
+ .iter()
+ .map(|x| x.as_list::<OffsetSize>())
+ .inspect(|l| {
+ output_len += l.len();
+ list_has_nulls |= l.null_count() != 0;
+ })
+ .collect::<Vec<_>>();
+
+ let mut dictionary_output_len = 0;
+ let dictionaries: Vec<_> = lists
+ .iter()
+ .map(|x| x.values().as_ref().as_dictionary::<K>())
+ .inspect(|d| dictionary_output_len += d.len())
+ .collect();
+
+ if !should_merge_dictionary_values::<K>(&dictionaries, dictionary_output_len) {
+ return concat_fallback(arrays, Capacities::Array(output_len));
+ }
+
+ let merged = merge_dictionary_values(&dictionaries, None)?;
+
+ let lists_nulls = list_has_nulls.then(|| {
+ let mut nulls = BooleanBufferBuilder::new(output_len);
+ for l in &lists {
+ match l.nulls() {
+ Some(n) => nulls.append_buffer(n.inner()),
+ None => nulls.append_n(l.len(), true),
+ }
+ }
+ NullBuffer::new(nulls.finish())
+ });
+
+ // Recompute keys
+ let mut key_values = Vec::with_capacity(dictionary_output_len);
+
+ let mut dictionary_has_nulls = false;
+ for (d, mapping) in dictionaries.iter().zip(merged.key_mappings) {
+ dictionary_has_nulls |= d.null_count() != 0;
+ for key in d.keys().values() {
+ // Use get to safely handle nulls
+ key_values.push(mapping.get(key.as_usize()).copied().unwrap_or_default())
+ }
+ }
+
+ let dictionary_nulls = dictionary_has_nulls.then(|| {
+ let mut nulls = BooleanBufferBuilder::new(dictionary_output_len);
+ for d in &dictionaries {
+ match d.nulls() {
+ Some(n) => nulls.append_buffer(n.inner()),
+ None => nulls.append_n(d.len(), true),
+ }
+ }
+ NullBuffer::new(nulls.finish())
+ });
+
+ let keys = PrimitiveArray::<K>::new(key_values.into(), dictionary_nulls);
+ // Sanity check
+ assert_eq!(keys.len(), dictionary_output_len);
+
+ let array = unsafe { DictionaryArray::new_unchecked(keys, merged.values) };
+
+ // Merge value offsets from the lists
+ let all_value_offsets_iterator = lists.iter().map(|x| x.offsets());
+
+ let value_offset_buffer = merge_value_offsets(all_value_offsets_iterator);
+
+ let builder = ArrayDataBuilder::new(arrays[0].data_type().clone())
+ .len(output_len)
+ .nulls(lists_nulls)
+ // `GenericListArray` must only have 1 buffer
+ .buffers(vec![value_offset_buffer])
+ // `GenericListArray` must only have 1 child_data
+ .child_data(vec![array.to_data()]);
+
+ // TODO - maybe use build_unchecked?
+ let array_data = builder.build()?;
+
+ let array = GenericListArray::<OffsetSize>::from(array_data);
+ Ok(Arc::new(array))
+}
+
+/// Merge value offsets
+///
+///
+/// if we have the following
+/// [[0, 3, 5], [0, 2, 2, 8], [], [0, 0, 1]]
+/// The output should be
+/// [ 0, 3, 5, 7, 7, 13, 13, 14]
+fn merge_value_offsets<
Review Comment:
should I move this to be in the offset buffer / scalar buffer / buffer file
##########
arrow-select/src/concat.rs:
##########
@@ -129,12 +130,164 @@ fn concat_dictionaries<K: ArrowDictionaryKeyType>(
Ok(Arc::new(array))
}
+fn concat_list_of_dictionaries<OffsetSize: OffsetSizeTrait, K: ArrowDictionaryKeyType>(
+ arrays: &[&dyn Array],
+) -> Result<ArrayRef, ArrowError> {
+ let mut output_len = 0;
+ let mut list_has_nulls = false;
+
+ let lists = arrays
+ .iter()
+ .map(|x| x.as_list::<OffsetSize>())
+ .inspect(|l| {
+ output_len += l.len();
+ list_has_nulls |= l.null_count() != 0;
+ })
+ .collect::<Vec<_>>();
+
+ let mut dictionary_output_len = 0;
+ let dictionaries: Vec<_> = lists
+ .iter()
+ .map(|x| x.values().as_ref().as_dictionary::<K>())
+ .inspect(|d| dictionary_output_len += d.len())
+ .collect();
+
+ if !should_merge_dictionary_values::<K>(&dictionaries, dictionary_output_len) {
+ return concat_fallback(arrays, Capacities::Array(output_len));
+ }
+
+ let merged = merge_dictionary_values(&dictionaries, None)?;
+
+ let lists_nulls = list_has_nulls.then(|| {
+ let mut nulls = BooleanBufferBuilder::new(output_len);
+ for l in &lists {
+ match l.nulls() {
+ Some(n) => nulls.append_buffer(n.inner()),
+ None => nulls.append_n(l.len(), true),
+ }
+ }
+ NullBuffer::new(nulls.finish())
+ });
+
+ // Recompute keys
+ let mut key_values = Vec::with_capacity(dictionary_output_len);
+
+ let mut dictionary_has_nulls = false;
+ for (d, mapping) in dictionaries.iter().zip(merged.key_mappings) {
+ dictionary_has_nulls |= d.null_count() != 0;
+ for key in d.keys().values() {
+ // Use get to safely handle nulls
+ key_values.push(mapping.get(key.as_usize()).copied().unwrap_or_default())
+ }
+ }
+
+ let dictionary_nulls = dictionary_has_nulls.then(|| {
+ let mut nulls = BooleanBufferBuilder::new(dictionary_output_len);
+ for d in &dictionaries {
+ match d.nulls() {
+ Some(n) => nulls.append_buffer(n.inner()),
+ None => nulls.append_n(d.len(), true),
+ }
+ }
+ NullBuffer::new(nulls.finish())
+ });
+
+ let keys = PrimitiveArray::<K>::new(key_values.into(), dictionary_nulls);
+ // Sanity check
+ assert_eq!(keys.len(), dictionary_output_len);
+
+ let array = unsafe { DictionaryArray::new_unchecked(keys, merged.values) };
+
+ // Merge value offsets from the lists
+ let all_value_offsets_iterator = lists.iter().map(|x| x.offsets());
+
+ let value_offset_buffer = merge_value_offsets(all_value_offsets_iterator);
+
+ let builder = ArrayDataBuilder::new(arrays[0].data_type().clone())
+ .len(output_len)
+ .nulls(lists_nulls)
+ // `GenericListArray` must only have 1 buffer
+ .buffers(vec![value_offset_buffer])
+ // `GenericListArray` must only have 1 child_data
+ .child_data(vec![array.to_data()]);
+
+ // TODO - maybe use build_unchecked?
+ let array_data = builder.build()?;
+
+ let array = GenericListArray::<OffsetSize>::from(array_data);
+ Ok(Arc::new(array))
+}
+
+/// Merge value offsets
+///
+///
+/// if we have the following
+/// [[0, 3, 5], [0, 2, 2, 8], [], [0, 0, 1]]
+/// The output should be
+/// [ 0, 3, 5, 7, 7, 13, 13, 14]
+fn merge_value_offsets<
Review Comment:
should I move this to be in the offset buffer / scalar buffer / buffer file?
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1890021873
##########
arrow-buffer/src/buffer/offset.rs:
##########
@@ -162,6 +163,44 @@ impl<O: ArrowNativeType> OffsetBuffer<O> {
}
}
+impl<O: ArrowNativeType + Integer + Copy> OffsetBuffer<O> {
Review Comment:
I would recommend a focus on getting something correct and working before trying to optimise it, and then when trying to optimise use empirical benchmarks.
FWIW I suspect this code will optimise very poorly as a result of the dynamic dispatch
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1894649607
##########
arrow-buffer/src/buffer/offset.rs:
##########
@@ -162,6 +163,44 @@ impl<O: ArrowNativeType> OffsetBuffer<O> {
}
}
+impl<O: ArrowNativeType + Integer + Copy> OffsetBuffer<O> {
Review Comment:
I'm fixing your comments now and I really like you approach! very readable and easy to use
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1932301566
##########
arrow-select/src/concat.rs:
##########
@@ -129,6 +129,54 @@ fn concat_dictionaries<K: ArrowDictionaryKeyType>(
Ok(Arc::new(array))
}
+fn concat_lists<OffsetSize: OffsetSizeTrait>(
+ arrays: &[&dyn Array],
+ field: &FieldRef,
+) -> Result<ArrayRef, ArrowError> {
+ let mut output_len = 0;
+ let mut list_has_nulls = false;
+
+ let lists = arrays
+ .iter()
+ .map(|x| x.as_list::<OffsetSize>())
+ .inspect(|l| {
+ output_len += l.len();
+ list_has_nulls |= l.null_count() != 0;
+ })
+ .collect::<Vec<_>>();
+
+ let lists_nulls = list_has_nulls.then(|| {
+ let mut nulls = BooleanBufferBuilder::new(output_len);
+ for l in &lists {
+ match l.nulls() {
+ Some(n) => nulls.append_buffer(n.inner()),
+ None => nulls.append_n(l.len(), true),
+ }
+ }
+ NullBuffer::new(nulls.finish())
+ });
+
+ let values: Vec<&dyn Array> = lists
+ .iter()
+ .map(|x| x.values().as_ref())
Review Comment:
I believe you are correct (also needs to slice values to get the last offset)
Proposed fix here:
- https://github.com/apache/arrow-rs/pull/7037
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#issuecomment-2619036256
Bug is related to concatenating sliced list arrays
- https://github.com/apache/arrow-rs/issues/7034
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1889289094
##########
arrow-select/src/concat.rs:
##########
@@ -169,6 +264,21 @@ pub fn concat(arrays: &[&dyn Array]) -> Result<ArrayRef, ArrowError> {
_ => unreachable!("illegal dictionary key type {k}")
};
} else {
+ if let DataType::List(field) = d {
+ if let DataType::Dictionary(k, _) = field.data_type() {
+ downcast_integer! {
+ k.as_ref() => (list_dict_helper, i32, arrays),
+ _ => unreachable!("illegal dictionary key type {k}")
+ };
+ }
+ } else if let DataType::LargeList(field) = d {
+ if let DataType::Dictionary(k, _) = field.data_type() {
+ downcast_integer! {
+ k.as_ref() => (list_dict_helper, i64, arrays),
+ _ => unreachable!("illegal dictionary key type {k}")
+ };
+ }
+ }
Review Comment:
This is quite a lot of codegen for what is a relatively niche use-case, I'm not sure we want to specialize 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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1889291117
##########
arrow-select/src/concat.rs:
##########
@@ -129,12 +130,106 @@ fn concat_dictionaries<K: ArrowDictionaryKeyType>(
Ok(Arc::new(array))
}
+fn concat_list_of_dictionaries<OffsetSize: OffsetSizeTrait, K: ArrowDictionaryKeyType>(
Review Comment:
It would significantly reduce the codegen to instead split concatenating the dictionary values, from re-encoding the offsets. As an added bonus this could also be done recursively.
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1889287132
##########
arrow-buffer/src/buffer/offset.rs:
##########
@@ -162,6 +163,44 @@ impl<O: ArrowNativeType> OffsetBuffer<O> {
}
}
+impl<O: ArrowNativeType + Integer + Copy> OffsetBuffer<O> {
+ /// Merge multiple [`OffsetBuffer`] into a single [`OffsetBuffer`]
+ ///
+ ///
+ /// ```
+ /// # use arrow_buffer::{OffsetBuffer, ScalarBuffer};
+ /// // [[0, 3, 5], [0, 2, 2, 8], [], [0, 0, 1]]
Review Comment:
FWIW `[]` is actually an illegal offset buffer
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1890021873
##########
arrow-buffer/src/buffer/offset.rs:
##########
@@ -162,6 +163,44 @@ impl<O: ArrowNativeType> OffsetBuffer<O> {
}
}
+impl<O: ArrowNativeType + Integer + Copy> OffsetBuffer<O> {
Review Comment:
I would recommend a focus on getting something correct and working before trying to optimise it, and then when trying to optimise use empirical benchmarks.
FWIW I suspect this code will optimise very poorly as a result of the dynamic dispatch.
TLDR I would recommend using from_lengths unless you can empirically show an alternative approach is sufficiently faster to warrant the additional complexity.
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1890029886
##########
arrow-buffer/src/buffer/offset.rs:
##########
@@ -162,6 +163,44 @@ impl<O: ArrowNativeType> OffsetBuffer<O> {
}
}
+impl<O: ArrowNativeType + Integer + Copy> OffsetBuffer<O> {
Review Comment:
Any branching, especially dynamic dispatch, will prevent LLVM from vectorising the loop
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1888719695
##########
arrow-select/src/concat.rs:
##########
@@ -129,12 +130,164 @@ fn concat_dictionaries<K: ArrowDictionaryKeyType>(
Ok(Arc::new(array))
}
+fn concat_list_of_dictionaries<OffsetSize: OffsetSizeTrait, K: ArrowDictionaryKeyType>(
+ arrays: &[&dyn Array],
+) -> Result<ArrayRef, ArrowError> {
+ let mut output_len = 0;
+ let mut list_has_nulls = false;
+
+ let lists = arrays
+ .iter()
+ .map(|x| x.as_list::<OffsetSize>())
+ .inspect(|l| {
+ output_len += l.len();
+ list_has_nulls |= l.null_count() != 0;
+ })
+ .collect::<Vec<_>>();
+
+ let mut dictionary_output_len = 0;
+ let dictionaries: Vec<_> = lists
+ .iter()
+ .map(|x| x.values().as_ref().as_dictionary::<K>())
+ .inspect(|d| dictionary_output_len += d.len())
+ .collect();
+
+ if !should_merge_dictionary_values::<K>(&dictionaries, dictionary_output_len) {
+ return concat_fallback(arrays, Capacities::Array(output_len));
+ }
+
+ let merged = merge_dictionary_values(&dictionaries, None)?;
+
+ let lists_nulls = list_has_nulls.then(|| {
+ let mut nulls = BooleanBufferBuilder::new(output_len);
+ for l in &lists {
+ match l.nulls() {
+ Some(n) => nulls.append_buffer(n.inner()),
+ None => nulls.append_n(l.len(), true),
+ }
+ }
+ NullBuffer::new(nulls.finish())
+ });
+
+ // Recompute keys
+ let mut key_values = Vec::with_capacity(dictionary_output_len);
+
+ let mut dictionary_has_nulls = false;
+ for (d, mapping) in dictionaries.iter().zip(merged.key_mappings) {
+ dictionary_has_nulls |= d.null_count() != 0;
+ for key in d.keys().values() {
+ // Use get to safely handle nulls
+ key_values.push(mapping.get(key.as_usize()).copied().unwrap_or_default())
+ }
+ }
+
+ let dictionary_nulls = dictionary_has_nulls.then(|| {
+ let mut nulls = BooleanBufferBuilder::new(dictionary_output_len);
+ for d in &dictionaries {
+ match d.nulls() {
+ Some(n) => nulls.append_buffer(n.inner()),
+ None => nulls.append_n(d.len(), true),
+ }
+ }
+ NullBuffer::new(nulls.finish())
+ });
+
+ let keys = PrimitiveArray::<K>::new(key_values.into(), dictionary_nulls);
+ // Sanity check
+ assert_eq!(keys.len(), dictionary_output_len);
+
+ let array = unsafe { DictionaryArray::new_unchecked(keys, merged.values) };
+
+ // Merge value offsets from the lists
+ let all_value_offsets_iterator = lists.iter().map(|x| x.offsets());
+
+ let value_offset_buffer = merge_value_offsets(all_value_offsets_iterator);
+
+ let builder = ArrayDataBuilder::new(arrays[0].data_type().clone())
+ .len(output_len)
+ .nulls(lists_nulls)
+ // `GenericListArray` must only have 1 buffer
+ .buffers(vec![value_offset_buffer])
+ // `GenericListArray` must only have 1 child_data
+ .child_data(vec![array.to_data()]);
+
+ // TODO - maybe use build_unchecked?
+ let array_data = builder.build()?;
+
+ let array = GenericListArray::<OffsetSize>::from(array_data);
+ Ok(Arc::new(array))
+}
+
+/// Merge value offsets
+///
+///
+/// if we have the following
+/// [[0, 3, 5], [0, 2, 2, 8], [], [0, 0, 1]]
+/// The output should be
+/// [ 0, 3, 5, 7, 7, 13, 13, 14]
+fn merge_value_offsets<
+ 'a,
+ OffsetSize: OffsetSizeTrait,
+ I: Iterator<Item = &'a OffsetBuffer<OffsetSize>>,
+>(
+ offset_buffers_iterator: I,
+) -> Buffer {
+ // 1. Filter out empty lists
+ let mut offset_buffers_iterator = offset_buffers_iterator.filter(|x| !x.is_empty());
+
+ // 2. Get first non-empty list as the starting point
+ let starting_buffer = offset_buffers_iterator.next();
+
+ // 3. If we have only empty lists, return an empty buffer
+ if starting_buffer.is_none() {
+ return Buffer::from(&[]);
+ }
+
+ let starting_buffer = starting_buffer.unwrap();
+
+ let mut offsets_iter: Box<dyn Iterator<Item = OffsetSize>> =
+ Box::new(starting_buffer.iter().copied());
+
+ // 4. Get the last value in the starting buffer as the starting point for the next buffer
+ // Safety: We already filtered out empty lists
+ let mut advance_by = *starting_buffer.last().unwrap();
+
+ // 5. Iterate over the remaining buffers
+ for offset_buffer in offset_buffers_iterator {
+ // 6. Get the last value of the current buffer so we can know how much to advance the next buffer
+ // Safety: We already filtered out empty lists
+ let last_value = *offset_buffer.last().unwrap();
+
+ // 7. Advance the offset buffer by the last value in the previous buffer
+ let offset_buffer_iter = offset_buffer
+ .iter()
+ // Skip the first value as it is the initial offset of 0
+ .skip(1)
+ .map(move |&x| x + advance_by);
+
+ // 8. concat the current buffer with the previous buffer
+ // Chaining keeps the iterator have trusting length
+ offsets_iter = Box::new(offsets_iter.chain(offset_buffer_iter));
Review Comment:
I have stack overflow error here, fixing now
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1894650417
##########
arrow-select/src/concat.rs:
##########
@@ -851,4 +1018,169 @@ mod tests {
assert_eq!(array.null_count(), 10);
assert_eq!(array.logical_null_count(), 10);
}
+
+ #[test]
+ fn concat_dictionary_list_array_simple() {
+ let scalars = vec![
+ create_single_row_list_of_dict(vec![Some("a")]),
+ create_single_row_list_of_dict(vec![Some("a")]),
+ create_single_row_list_of_dict(vec![Some("b")]),
+ ];
+
+ let arrays = scalars.iter().map(|a| a as &(dyn Array)).collect::<Vec<_>>();
+ let concat_res = concat(arrays.as_slice()).unwrap();
+
+ let expected_list = create_list_of_dict(vec![
+ // Row 1
+ Some(vec![Some("a")]),
+ Some(vec![Some("a")]),
+ Some(vec![Some("b")]),
+ ]);
+
+ let list = concat_res.as_list::<i32>();
+
+ // Assert that the list is equal to the expected list
+ list.iter().zip(expected_list.iter()).for_each(|(a, b)| {
+ assert_eq!(a, b);
+ });
+
+ let dict = list
+ .values()
+ .as_dictionary::<Int32Type>()
+ .downcast_dict::<StringArray>()
+ .unwrap();
+ println!("{:?}", dict);
+
+ assert_dictionary_has_unique_values::<_, StringArray>(
+ list.values().as_dictionary::<Int32Type>(),
+ );
+ }
+
+ #[test]
+ fn concat_dictionary_list_array_with_multiple_rows() {
Review Comment:
removed
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(arrow-select): `concat` kernel will merge dictionary values for list of dictionaries [arrow-rs]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on code in PR #6893:
URL: https://github.com/apache/arrow-rs/pull/6893#discussion_r1894655245
##########
arrow-select/src/concat.rs:
##########
@@ -129,12 +131,161 @@ fn concat_dictionaries<K: ArrowDictionaryKeyType>(
Ok(Arc::new(array))
}
+fn concat_list_of_dictionaries<OffsetSize: OffsetSizeTrait, K: ArrowDictionaryKeyType>(
+ arrays: &[&dyn Array],
+) -> Result<ArrayRef, ArrowError> {
+ let mut output_len = 0;
+ let mut list_has_nulls = false;
+
+ let lists = arrays
+ .iter()
+ .map(|x| x.as_list::<OffsetSize>())
+ .inspect(|l| {
+ output_len += l.len();
+ list_has_nulls |= l.null_count() != 0;
+ })
+ .collect::<Vec<_>>();
+
+ let mut dictionary_output_len = 0;
+ let dictionaries: Vec<_> = lists
+ .iter()
+ .map(|x| x.values().as_ref().as_dictionary::<K>())
+ .inspect(|d| dictionary_output_len += d.len())
+ .collect();
+
+ if !should_merge_dictionary_values::<K>(&dictionaries, dictionary_output_len) {
+ return concat_fallback(arrays, Capacities::Array(output_len));
+ }
+
+ let merged = merge_dictionary_values(&dictionaries, None)?;
+
+ let lists_nulls = list_has_nulls.then(|| {
+ let mut nulls = BooleanBufferBuilder::new(output_len);
+ for l in &lists {
+ match l.nulls() {
+ Some(n) => nulls.append_buffer(n.inner()),
+ None => nulls.append_n(l.len(), true),
+ }
+ }
+ NullBuffer::new(nulls.finish())
+ });
+
+ // Recompute keys
+ let mut key_values = Vec::with_capacity(dictionary_output_len);
+
+ let mut dictionary_has_nulls = false;
+ for (d, mapping) in dictionaries.iter().zip(merged.key_mappings) {
+ dictionary_has_nulls |= d.null_count() != 0;
+ for key in d.keys().values() {
+ // Use get to safely handle nulls
+ key_values.push(mapping.get(key.as_usize()).copied().unwrap_or_default())
+ }
+ }
+
+ let dictionary_nulls = dictionary_has_nulls.then(|| {
+ let mut nulls = BooleanBufferBuilder::new(dictionary_output_len);
+ for d in &dictionaries {
+ match d.nulls() {
+ Some(n) => nulls.append_buffer(n.inner()),
+ None => nulls.append_n(d.len(), true),
+ }
+ }
+ NullBuffer::new(nulls.finish())
+ });
+
+ let keys = PrimitiveArray::<K>::new(key_values.into(), dictionary_nulls);
+ // Sanity check
+ assert_eq!(keys.len(), dictionary_output_len);
+
+ let array = unsafe { DictionaryArray::new_unchecked(keys, merged.values) };
+
+ // Merge value offsets from the lists
+ let all_value_offsets_iterator = lists
+ .iter()
+ .map(|x| x.offsets());
+
+ let value_offset_buffer = merge_value_offsets(all_value_offsets_iterator);
+
+ let builder = ArrayDataBuilder::new(arrays[0].data_type().clone())
+ .len(output_len)
+ .nulls(lists_nulls)
+ // `GenericListArray` must only have 1 buffer
+ .buffers(vec![value_offset_buffer])
+ // `GenericListArray` must only have 1 child_data
+ .child_data(vec![array.to_data()]);
+
+ // TODO - maybe use build_unchecked?
+ let array_data = builder.build()?;
Review Comment:
@tustvold wdyt?
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org