You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "tustvold (via GitHub)" <gi...@apache.org> on 2023/06/13 11:13:46 UTC
[GitHub] [arrow-rs] tustvold opened a new pull request, #4405: Improve take primitive performance (#4404)
tustvold opened a new pull request, #4405:
URL: https://github.com/apache/arrow-rs/pull/4405
# Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123.
-->
Closes #4404
# Rationale for this change
<!--
Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes.
-->
Fixes #4404 and improves performance significantly
Using the benchmarks in #4403
```
take i32 512 time: [259.96 ns 260.19 ns 260.43 ns]
change: [-33.771% -33.654% -33.538%] (p = 0.00 < 0.05)
Performance has improved.
Found 4 outliers among 100 measurements (4.00%)
2 (2.00%) high mild
2 (2.00%) high severe
take i32 1024 time: [411.64 ns 411.83 ns 412.05 ns]
change: [-31.098% -31.057% -31.012%] (p = 0.00 < 0.05)
Performance has improved.
Found 2 outliers among 100 measurements (2.00%)
2 (2.00%) high mild
take i32 null indices 1024
time: [556.36 ns 556.65 ns 556.94 ns]
change: [-14.529% -14.467% -14.407%] (p = 0.00 < 0.05)
Performance has improved.
take i32 null values 1024
time: [1.3158 µs 1.3167 µs 1.3176 µs]
change: [-41.607% -41.530% -41.446%] (p = 0.00 < 0.05)
Performance has improved.
Found 3 outliers among 100 measurements (3.00%)
2 (2.00%) high mild
1 (1.00%) high severe
take i32 null values null indices 1024
time: [1.6743 µs 1.6761 µs 1.6775 µs]
change: [-43.953% -43.710% -43.467%] (p = 0.00 < 0.05)
Performance has improved.
take check bounds i32 512
time: [381.76 ns 381.99 ns 382.22 ns]
change: [-25.746% -25.670% -25.599%] (p = 0.00 < 0.05)
Performance has improved.
Found 3 outliers among 100 measurements (3.00%)
2 (2.00%) high mild
1 (1.00%) high severe
take check bounds i32 1024
time: [662.78 ns 663.02 ns 663.26 ns]
change: [-22.074% -21.942% -21.848%] (p = 0.00 < 0.05)
Performance has improved.
take bool 512 time: [525.34 ns 525.72 ns 526.11 ns]
change: [-6.8248% -6.5847% -6.3665%] (p = 0.00 < 0.05)
Performance has improved.
Found 5 outliers among 100 measurements (5.00%)
3 (3.00%) low severe
1 (1.00%) high mild
1 (1.00%) high severe
take bool 1024 time: [885.39 ns 886.30 ns 887.29 ns]
change: [-12.935% -12.734% -12.536%] (p = 0.00 < 0.05)
Performance has improved.
Found 5 outliers among 100 measurements (5.00%)
5 (5.00%) high mild
take bool null indices 1024
time: [1.0396 µs 1.0401 µs 1.0406 µs]
change: [-51.963% -51.923% -51.885%] (p = 0.00 < 0.05)
Performance has improved.
Found 5 outliers among 100 measurements (5.00%)
3 (3.00%) high mild
2 (2.00%) high severe
take bool null values 1024
time: [1.7840 µs 1.7853 µs 1.7871 µs]
change: [-2.8921% -2.7540% -2.6122%] (p = 0.00 < 0.05)
Performance has improved.
Found 5 outliers among 100 measurements (5.00%)
3 (3.00%) high mild
2 (2.00%) high severe
take bool null values null indices 1024
time: [2.1861 µs 2.1878 µs 2.1897 µs]
change: [-45.707% -45.587% -45.456%] (p = 0.00 < 0.05)
Performance has improved.
Found 1 outliers among 100 measurements (1.00%)
```
# What changes are included in this PR?
<!--
There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR.
-->
# Are there any user-facing changes?
<!--
If there are user-facing changes then we may require documentation to be updated before approving the PR.
-->
<!---
If there are any breaking changes to public APIs, please add the `breaking change` label.
-->
--
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
[GitHub] [arrow-rs] tustvold commented on a diff in pull request #4405: Improve take primitive performance (#4404)
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #4405:
URL: https://github.com/apache/arrow-rs/pull/4405#discussion_r1228118977
##########
arrow-select/src/take.rs:
##########
@@ -374,148 +243,91 @@ fn take_primitive<T, I>(
where
T: ArrowPrimitiveType,
I: ArrowPrimitiveType,
- I::Native: ToPrimitive,
{
- let indices_nulls = indices.nulls().filter(|x| x.null_count() > 0);
- let values_nulls = values.nulls().filter(|x| x.null_count() > 0);
-
- // note: this function should only panic when "an index is not null and out of bounds".
- // if the index is null, its value is undefined and therefore we should not read from it.
- let (buffer, nulls) = match (values_nulls, indices_nulls) {
- (None, None) => {
- // * no nulls
- // * all `indices.values()` are valid
- take_no_nulls(values.values(), indices.values())?
- }
- (Some(values_nulls), None) => {
- // * nulls come from `values` alone
- // * all `indices.values()` are valid
- take_values_nulls(values.values(), values_nulls, indices.values())?
- }
- (None, Some(indices_nulls)) => {
- // in this branch it is unsound to read and use `index.values()`,
- // as doing so is UB when they come from a null slot.
- take_indices_nulls(values.values(), indices.values(), indices_nulls)?
- }
- (Some(values_nulls), Some(indices_nulls)) => {
- // in this branch it is unsound to read and use `index.values()`,
- // as doing so is UB when they come from a null slot.
- take_values_indices_nulls(
- values.values(),
- values_nulls,
- indices.values(),
- indices_nulls,
- )?
- }
- };
-
- let data = unsafe {
- ArrayData::new_unchecked(
- values.data_type().clone(),
- indices.len(),
- None,
- nulls,
- 0,
- vec![buffer],
- vec![],
- )
- };
- Ok(PrimitiveArray::<T>::from(data))
+ let values_buf = take_native(values.values(), indices);
+ let nulls = take_nulls(values.nulls(), indices);
+ Ok(PrimitiveArray::new(values_buf, nulls).with_data_type(values.data_type().clone()))
}
-fn take_bits<IndexType>(
- values: &Buffer,
- values_offset: usize,
- indices: &PrimitiveArray<IndexType>,
-) -> Result<Buffer, ArrowError>
-where
- IndexType: ArrowPrimitiveType,
- IndexType::Native: ToPrimitive,
-{
- let len = indices.len();
- let values_slice = values.as_slice();
- let mut output_buffer = MutableBuffer::new_null(len);
- let output_slice = output_buffer.as_slice_mut();
-
- let indices_has_nulls = indices.null_count() > 0;
+#[inline(never)]
+fn take_nulls<I: ArrowPrimitiveType>(
+ values: Option<&NullBuffer>,
+ indices: &PrimitiveArray<I>,
+) -> Option<NullBuffer> {
+ match values.filter(|n| n.null_count() > 0) {
+ Some(n) => {
+ let buffer = take_bits(n.inner(), indices);
+ Some(NullBuffer::new(buffer)).filter(|n| n.null_count() > 0)
+ }
+ None => indices.nulls().cloned(),
+ }
+}
- if indices_has_nulls {
- indices
+#[inline(never)]
+fn take_native<T: ArrowNativeType, I: ArrowPrimitiveType>(
+ values: &[T],
+ indices: &PrimitiveArray<I>,
+) -> ScalarBuffer<T> {
+ match indices.nulls().filter(|n| n.null_count() > 0) {
+ Some(n) => indices
+ .values()
.iter()
.enumerate()
- .try_for_each::<_, Result<(), ArrowError>>(|(i, index)| {
- if let Some(index) = index {
- let index = ToPrimitive::to_usize(&index).ok_or_else(|| {
- ArrowError::ComputeError("Cast to usize failed".to_string())
- })?;
-
- if bit_util::get_bit(values_slice, values_offset + index) {
- bit_util::set_bit(output_slice, i);
- }
- }
-
- Ok(())
- })?;
- } else {
- indices
+ .map(|(idx, index)| match values.get(index.as_usize()) {
+ Some(v) => *v,
+ None => match n.is_null(idx) {
+ true => T::default(),
+ false => panic!("Out-of-bounds index {index:?}"),
+ },
+ })
+ .collect(),
+ None => indices
.values()
.iter()
- .enumerate()
- .try_for_each::<_, Result<(), ArrowError>>(|(i, index)| {
- let index = ToPrimitive::to_usize(index).ok_or_else(|| {
- ArrowError::ComputeError("Cast to usize failed".to_string())
- })?;
+ .map(|index| values[index.as_usize()])
+ .collect(),
+ }
+}
- if bit_util::get_bit(values_slice, values_offset + index) {
- bit_util::set_bit(output_slice, i);
- }
- Ok(())
- })?;
+#[inline(never)]
+fn take_bits<I: ArrowPrimitiveType>(
+ values: &BooleanBuffer,
+ indices: &PrimitiveArray<I>,
+) -> BooleanBuffer {
+ let len = indices.len();
+ let mut output_buffer = MutableBuffer::new_null(len);
+ let output_slice = output_buffer.as_slice_mut();
+
+ match indices.nulls().filter(|n| n.null_count() > 0) {
+ Some(nulls) => nulls.valid_indices().for_each(|idx| {
+ if values.value(indices.value(idx).as_usize()) {
+ bit_util::set_bit(output_slice, idx);
+ }
+ }),
+ None => indices.values().iter().enumerate().for_each(|(i, index)| {
+ if values.value(index.as_usize()) {
Review Comment:
In theory BooleanBuffer::collect_bool should be faster, in this case it turned out to be slower for some reason - likely something to do with what LLVM is doing with the bound checks
--
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
[GitHub] [arrow-rs] Dandandan commented on a diff in pull request #4405: Improve `take` kernel performance on primitive arrays, fix bad null index handling (#4404)
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on code in PR #4405:
URL: https://github.com/apache/arrow-rs/pull/4405#discussion_r1228154218
##########
arrow-select/src/take.rs:
##########
@@ -374,148 +243,91 @@ fn take_primitive<T, I>(
where
T: ArrowPrimitiveType,
I: ArrowPrimitiveType,
- I::Native: ToPrimitive,
{
- let indices_nulls = indices.nulls().filter(|x| x.null_count() > 0);
- let values_nulls = values.nulls().filter(|x| x.null_count() > 0);
-
- // note: this function should only panic when "an index is not null and out of bounds".
- // if the index is null, its value is undefined and therefore we should not read from it.
- let (buffer, nulls) = match (values_nulls, indices_nulls) {
- (None, None) => {
- // * no nulls
- // * all `indices.values()` are valid
- take_no_nulls(values.values(), indices.values())?
- }
- (Some(values_nulls), None) => {
- // * nulls come from `values` alone
- // * all `indices.values()` are valid
- take_values_nulls(values.values(), values_nulls, indices.values())?
- }
- (None, Some(indices_nulls)) => {
- // in this branch it is unsound to read and use `index.values()`,
- // as doing so is UB when they come from a null slot.
- take_indices_nulls(values.values(), indices.values(), indices_nulls)?
- }
- (Some(values_nulls), Some(indices_nulls)) => {
- // in this branch it is unsound to read and use `index.values()`,
- // as doing so is UB when they come from a null slot.
- take_values_indices_nulls(
- values.values(),
- values_nulls,
- indices.values(),
- indices_nulls,
- )?
- }
- };
-
- let data = unsafe {
- ArrayData::new_unchecked(
- values.data_type().clone(),
- indices.len(),
- None,
- nulls,
- 0,
- vec![buffer],
- vec![],
- )
- };
- Ok(PrimitiveArray::<T>::from(data))
+ let values_buf = take_native(values.values(), indices);
+ let nulls = take_nulls(values.nulls(), indices);
+ Ok(PrimitiveArray::new(values_buf, nulls).with_data_type(values.data_type().clone()))
}
-fn take_bits<IndexType>(
- values: &Buffer,
- values_offset: usize,
- indices: &PrimitiveArray<IndexType>,
-) -> Result<Buffer, ArrowError>
-where
- IndexType: ArrowPrimitiveType,
- IndexType::Native: ToPrimitive,
-{
- let len = indices.len();
- let values_slice = values.as_slice();
- let mut output_buffer = MutableBuffer::new_null(len);
- let output_slice = output_buffer.as_slice_mut();
-
- let indices_has_nulls = indices.null_count() > 0;
+#[inline(never)]
+fn take_nulls<I: ArrowPrimitiveType>(
+ values: Option<&NullBuffer>,
+ indices: &PrimitiveArray<I>,
+) -> Option<NullBuffer> {
+ match values.filter(|n| n.null_count() > 0) {
+ Some(n) => {
+ let buffer = take_bits(n.inner(), indices);
+ Some(NullBuffer::new(buffer)).filter(|n| n.null_count() > 0)
+ }
+ None => indices.nulls().cloned(),
+ }
+}
- if indices_has_nulls {
- indices
+#[inline(never)]
+fn take_native<T: ArrowNativeType, I: ArrowPrimitiveType>(
+ values: &[T],
+ indices: &PrimitiveArray<I>,
+) -> ScalarBuffer<T> {
+ match indices.nulls().filter(|n| n.null_count() > 0) {
+ Some(n) => indices
+ .values()
.iter()
.enumerate()
- .try_for_each::<_, Result<(), ArrowError>>(|(i, index)| {
- if let Some(index) = index {
- let index = ToPrimitive::to_usize(&index).ok_or_else(|| {
- ArrowError::ComputeError("Cast to usize failed".to_string())
- })?;
-
- if bit_util::get_bit(values_slice, values_offset + index) {
- bit_util::set_bit(output_slice, i);
- }
- }
-
- Ok(())
- })?;
- } else {
- indices
+ .map(|(idx, index)| match values.get(index.as_usize()) {
+ Some(v) => *v,
+ None => match n.is_null(idx) {
+ true => T::default(),
+ false => panic!("Out-of-bounds index {index:?}"),
+ },
+ })
+ .collect(),
+ None => indices
.values()
.iter()
- .enumerate()
- .try_for_each::<_, Result<(), ArrowError>>(|(i, index)| {
- let index = ToPrimitive::to_usize(index).ok_or_else(|| {
- ArrowError::ComputeError("Cast to usize failed".to_string())
- })?;
+ .map(|index| values[index.as_usize()])
+ .collect(),
+ }
+}
- if bit_util::get_bit(values_slice, values_offset + index) {
- bit_util::set_bit(output_slice, i);
- }
- Ok(())
- })?;
+#[inline(never)]
+fn take_bits<I: ArrowPrimitiveType>(
+ values: &BooleanBuffer,
+ indices: &PrimitiveArray<I>,
+) -> BooleanBuffer {
+ let len = indices.len();
+ let mut output_buffer = MutableBuffer::new_null(len);
+ let output_slice = output_buffer.as_slice_mut();
+
+ match indices.nulls().filter(|n| n.null_count() > 0) {
+ Some(nulls) => nulls.valid_indices().for_each(|idx| {
+ if values.value(indices.value(idx).as_usize()) {
+ bit_util::set_bit(output_slice, idx);
+ }
+ }),
+ None => indices.values().iter().enumerate().for_each(|(i, index)| {
+ if values.value(index.as_usize()) {
Review Comment:
ok, thanks @tustvold
--
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
[GitHub] [arrow-rs] tustvold commented on a diff in pull request #4405: Improve take primitive performance (#4404)
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #4405:
URL: https://github.com/apache/arrow-rs/pull/4405#discussion_r1227954866
##########
arrow-buffer/src/buffer/scalar.rs:
##########
@@ -140,6 +140,12 @@ impl<T: ArrowNativeType> From<Vec<T>> for ScalarBuffer<T> {
}
}
+impl<T: ArrowNativeType> FromIterator<T> for ScalarBuffer<T> {
+ fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
+ iter.into_iter().collect::<Vec<_>>().into()
Review Comment:
An important thing to note is that `Vec: FromIterator` has a specialization for `TrustedLen` iterators, such as those from slices. This allows us to not need `Buffer::try_from_trusted_len_iter`
--
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
[GitHub] [arrow-rs] tustvold commented on a diff in pull request #4405: Improve take primitive performance (#4404)
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #4405:
URL: https://github.com/apache/arrow-rs/pull/4405#discussion_r1227953255
##########
arrow-buffer/src/buffer/boolean.rs:
##########
@@ -128,6 +128,7 @@ impl BooleanBuffer {
/// # Panics
///
/// Panics if `i >= self.len()`
+ #[inline]
Review Comment:
This is vital for the performance of take_bits
--
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
[GitHub] [arrow-rs] tustvold commented on a diff in pull request #4405: Improve take primitive performance (#4404)
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #4405:
URL: https://github.com/apache/arrow-rs/pull/4405#discussion_r1227953969
##########
arrow-select/src/take.rs:
##########
@@ -2155,4 +1977,19 @@ mod tests {
UInt32Array::from(vec![9, 10, 11, 6, 7, 8, 3, 4, 5, 6, 7, 8, 0, 1, 2])
);
}
+
+ #[test]
+ fn test_take_null_indices() {
Review Comment:
Test for #4404
--
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
[GitHub] [arrow-rs] tustvold commented on a diff in pull request #4405: Improve take primitive performance (#4404)
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #4405:
URL: https://github.com/apache/arrow-rs/pull/4405#discussion_r1227953610
##########
arrow-select/src/take.rs:
##########
@@ -376,134 +253,87 @@ where
I: ArrowPrimitiveType,
I::Native: ToPrimitive,
{
- let indices_nulls = indices.nulls().filter(|x| x.null_count() > 0);
- let values_nulls = values.nulls().filter(|x| x.null_count() > 0);
-
- // note: this function should only panic when "an index is not null and out of bounds".
- // if the index is null, its value is undefined and therefore we should not read from it.
- let (buffer, nulls) = match (values_nulls, indices_nulls) {
- (None, None) => {
- // * no nulls
- // * all `indices.values()` are valid
- take_no_nulls(values.values(), indices.values())?
- }
- (Some(values_nulls), None) => {
- // * nulls come from `values` alone
- // * all `indices.values()` are valid
- take_values_nulls(values.values(), values_nulls, indices.values())?
- }
- (None, Some(indices_nulls)) => {
- // in this branch it is unsound to read and use `index.values()`,
- // as doing so is UB when they come from a null slot.
- take_indices_nulls(values.values(), indices.values(), indices_nulls)?
- }
- (Some(values_nulls), Some(indices_nulls)) => {
- // in this branch it is unsound to read and use `index.values()`,
- // as doing so is UB when they come from a null slot.
- take_values_indices_nulls(
- values.values(),
- values_nulls,
- indices.values(),
- indices_nulls,
- )?
- }
- };
-
- let data = unsafe {
- ArrayData::new_unchecked(
- values.data_type().clone(),
- indices.len(),
- None,
- nulls,
- 0,
- vec![buffer],
- vec![],
- )
- };
- Ok(PrimitiveArray::<T>::from(data))
+ let values_buf = take_native(values.values(), indices);
+ let nulls = take_nulls(values.nulls(), indices);
+ Ok(PrimitiveArray::new(values_buf, nulls).with_data_type(values.data_type().clone()))
}
-fn take_bits<IndexType>(
- values: &Buffer,
- values_offset: usize,
- indices: &PrimitiveArray<IndexType>,
-) -> Result<Buffer, ArrowError>
-where
- IndexType: ArrowPrimitiveType,
- IndexType::Native: ToPrimitive,
-{
- let len = indices.len();
- let values_slice = values.as_slice();
- let mut output_buffer = MutableBuffer::new_null(len);
- let output_slice = output_buffer.as_slice_mut();
-
- let indices_has_nulls = indices.null_count() > 0;
+#[inline(never)]
Review Comment:
I had issues with LLVM inlining and then not optimising correctly, this just forces it to not be stupid
--
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
[GitHub] [arrow-rs] Dandandan commented on a diff in pull request #4405: Improve take primitive performance (#4404)
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on code in PR #4405:
URL: https://github.com/apache/arrow-rs/pull/4405#discussion_r1228123182
##########
arrow-select/src/take.rs:
##########
@@ -374,148 +243,91 @@ fn take_primitive<T, I>(
where
T: ArrowPrimitiveType,
I: ArrowPrimitiveType,
- I::Native: ToPrimitive,
{
- let indices_nulls = indices.nulls().filter(|x| x.null_count() > 0);
- let values_nulls = values.nulls().filter(|x| x.null_count() > 0);
-
- // note: this function should only panic when "an index is not null and out of bounds".
- // if the index is null, its value is undefined and therefore we should not read from it.
- let (buffer, nulls) = match (values_nulls, indices_nulls) {
- (None, None) => {
- // * no nulls
- // * all `indices.values()` are valid
- take_no_nulls(values.values(), indices.values())?
- }
- (Some(values_nulls), None) => {
- // * nulls come from `values` alone
- // * all `indices.values()` are valid
- take_values_nulls(values.values(), values_nulls, indices.values())?
- }
- (None, Some(indices_nulls)) => {
- // in this branch it is unsound to read and use `index.values()`,
- // as doing so is UB when they come from a null slot.
- take_indices_nulls(values.values(), indices.values(), indices_nulls)?
- }
- (Some(values_nulls), Some(indices_nulls)) => {
- // in this branch it is unsound to read and use `index.values()`,
- // as doing so is UB when they come from a null slot.
- take_values_indices_nulls(
- values.values(),
- values_nulls,
- indices.values(),
- indices_nulls,
- )?
- }
- };
-
- let data = unsafe {
- ArrayData::new_unchecked(
- values.data_type().clone(),
- indices.len(),
- None,
- nulls,
- 0,
- vec![buffer],
- vec![],
- )
- };
- Ok(PrimitiveArray::<T>::from(data))
+ let values_buf = take_native(values.values(), indices);
+ let nulls = take_nulls(values.nulls(), indices);
+ Ok(PrimitiveArray::new(values_buf, nulls).with_data_type(values.data_type().clone()))
}
-fn take_bits<IndexType>(
- values: &Buffer,
- values_offset: usize,
- indices: &PrimitiveArray<IndexType>,
-) -> Result<Buffer, ArrowError>
-where
- IndexType: ArrowPrimitiveType,
- IndexType::Native: ToPrimitive,
-{
- let len = indices.len();
- let values_slice = values.as_slice();
- let mut output_buffer = MutableBuffer::new_null(len);
- let output_slice = output_buffer.as_slice_mut();
-
- let indices_has_nulls = indices.null_count() > 0;
+#[inline(never)]
+fn take_nulls<I: ArrowPrimitiveType>(
+ values: Option<&NullBuffer>,
+ indices: &PrimitiveArray<I>,
+) -> Option<NullBuffer> {
+ match values.filter(|n| n.null_count() > 0) {
+ Some(n) => {
+ let buffer = take_bits(n.inner(), indices);
+ Some(NullBuffer::new(buffer)).filter(|n| n.null_count() > 0)
+ }
+ None => indices.nulls().cloned(),
+ }
+}
- if indices_has_nulls {
- indices
+#[inline(never)]
+fn take_native<T: ArrowNativeType, I: ArrowPrimitiveType>(
+ values: &[T],
+ indices: &PrimitiveArray<I>,
+) -> ScalarBuffer<T> {
+ match indices.nulls().filter(|n| n.null_count() > 0) {
+ Some(n) => indices
+ .values()
.iter()
.enumerate()
- .try_for_each::<_, Result<(), ArrowError>>(|(i, index)| {
- if let Some(index) = index {
- let index = ToPrimitive::to_usize(&index).ok_or_else(|| {
- ArrowError::ComputeError("Cast to usize failed".to_string())
- })?;
-
- if bit_util::get_bit(values_slice, values_offset + index) {
- bit_util::set_bit(output_slice, i);
- }
- }
-
- Ok(())
- })?;
- } else {
- indices
+ .map(|(idx, index)| match values.get(index.as_usize()) {
+ Some(v) => *v,
+ None => match n.is_null(idx) {
+ true => T::default(),
+ false => panic!("Out-of-bounds index {index:?}"),
+ },
+ })
+ .collect(),
+ None => indices
.values()
.iter()
- .enumerate()
- .try_for_each::<_, Result<(), ArrowError>>(|(i, index)| {
- let index = ToPrimitive::to_usize(index).ok_or_else(|| {
- ArrowError::ComputeError("Cast to usize failed".to_string())
- })?;
+ .map(|index| values[index.as_usize()])
+ .collect(),
+ }
+}
- if bit_util::get_bit(values_slice, values_offset + index) {
- bit_util::set_bit(output_slice, i);
- }
- Ok(())
- })?;
+#[inline(never)]
+fn take_bits<I: ArrowPrimitiveType>(
+ values: &BooleanBuffer,
+ indices: &PrimitiveArray<I>,
+) -> BooleanBuffer {
+ let len = indices.len();
+ let mut output_buffer = MutableBuffer::new_null(len);
+ let output_slice = output_buffer.as_slice_mut();
+
+ match indices.nulls().filter(|n| n.null_count() > 0) {
+ Some(nulls) => nulls.valid_indices().for_each(|idx| {
+ if values.value(indices.value(idx).as_usize()) {
+ bit_util::set_bit(output_slice, idx);
+ }
+ }),
+ None => indices.values().iter().enumerate().for_each(|(i, index)| {
+ if values.value(index.as_usize()) {
Review Comment:
Yeah `collect_bool`
--
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
[GitHub] [arrow-rs] alamb commented on a diff in pull request #4405: Improve `take` kernel performance on primitive arrays, fix bad null index handling (#4404)
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #4405:
URL: https://github.com/apache/arrow-rs/pull/4405#discussion_r1228144351
##########
arrow-select/src/take.rs:
##########
@@ -374,148 +243,91 @@ fn take_primitive<T, I>(
where
T: ArrowPrimitiveType,
I: ArrowPrimitiveType,
- I::Native: ToPrimitive,
{
- let indices_nulls = indices.nulls().filter(|x| x.null_count() > 0);
- let values_nulls = values.nulls().filter(|x| x.null_count() > 0);
-
- // note: this function should only panic when "an index is not null and out of bounds".
- // if the index is null, its value is undefined and therefore we should not read from it.
- let (buffer, nulls) = match (values_nulls, indices_nulls) {
- (None, None) => {
- // * no nulls
- // * all `indices.values()` are valid
- take_no_nulls(values.values(), indices.values())?
- }
- (Some(values_nulls), None) => {
- // * nulls come from `values` alone
- // * all `indices.values()` are valid
- take_values_nulls(values.values(), values_nulls, indices.values())?
- }
- (None, Some(indices_nulls)) => {
- // in this branch it is unsound to read and use `index.values()`,
- // as doing so is UB when they come from a null slot.
- take_indices_nulls(values.values(), indices.values(), indices_nulls)?
- }
- (Some(values_nulls), Some(indices_nulls)) => {
- // in this branch it is unsound to read and use `index.values()`,
- // as doing so is UB when they come from a null slot.
- take_values_indices_nulls(
- values.values(),
- values_nulls,
- indices.values(),
- indices_nulls,
- )?
- }
- };
-
- let data = unsafe {
- ArrayData::new_unchecked(
- values.data_type().clone(),
- indices.len(),
- None,
- nulls,
- 0,
- vec![buffer],
- vec![],
- )
- };
- Ok(PrimitiveArray::<T>::from(data))
+ let values_buf = take_native(values.values(), indices);
+ let nulls = take_nulls(values.nulls(), indices);
+ Ok(PrimitiveArray::new(values_buf, nulls).with_data_type(values.data_type().clone()))
}
-fn take_bits<IndexType>(
- values: &Buffer,
- values_offset: usize,
- indices: &PrimitiveArray<IndexType>,
-) -> Result<Buffer, ArrowError>
-where
- IndexType: ArrowPrimitiveType,
- IndexType::Native: ToPrimitive,
-{
- let len = indices.len();
- let values_slice = values.as_slice();
- let mut output_buffer = MutableBuffer::new_null(len);
- let output_slice = output_buffer.as_slice_mut();
-
- let indices_has_nulls = indices.null_count() > 0;
+#[inline(never)]
+fn take_nulls<I: ArrowPrimitiveType>(
+ values: Option<&NullBuffer>,
+ indices: &PrimitiveArray<I>,
+) -> Option<NullBuffer> {
+ match values.filter(|n| n.null_count() > 0) {
Review Comment:
it is certainly neat to see nice Rust code like this and then know `rustc` / `LLVM` did the right thing to make it fast
##########
arrow-buffer/src/buffer/scalar.rs:
##########
@@ -140,6 +140,12 @@ impl<T: ArrowNativeType> From<Vec<T>> for ScalarBuffer<T> {
}
}
+impl<T: ArrowNativeType> FromIterator<T> for ScalarBuffer<T> {
+ fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
+ iter.into_iter().collect::<Vec<_>>().into()
Review Comment:
Can you please add this (very interesting) information as a comment inline?
##########
arrow-select/src/take.rs:
##########
@@ -376,134 +253,87 @@ where
I: ArrowPrimitiveType,
I::Native: ToPrimitive,
{
- let indices_nulls = indices.nulls().filter(|x| x.null_count() > 0);
- let values_nulls = values.nulls().filter(|x| x.null_count() > 0);
-
- // note: this function should only panic when "an index is not null and out of bounds".
- // if the index is null, its value is undefined and therefore we should not read from it.
- let (buffer, nulls) = match (values_nulls, indices_nulls) {
- (None, None) => {
- // * no nulls
- // * all `indices.values()` are valid
- take_no_nulls(values.values(), indices.values())?
- }
- (Some(values_nulls), None) => {
- // * nulls come from `values` alone
- // * all `indices.values()` are valid
- take_values_nulls(values.values(), values_nulls, indices.values())?
- }
- (None, Some(indices_nulls)) => {
- // in this branch it is unsound to read and use `index.values()`,
- // as doing so is UB when they come from a null slot.
- take_indices_nulls(values.values(), indices.values(), indices_nulls)?
- }
- (Some(values_nulls), Some(indices_nulls)) => {
- // in this branch it is unsound to read and use `index.values()`,
- // as doing so is UB when they come from a null slot.
- take_values_indices_nulls(
- values.values(),
- values_nulls,
- indices.values(),
- indices_nulls,
- )?
- }
- };
-
- let data = unsafe {
- ArrayData::new_unchecked(
- values.data_type().clone(),
- indices.len(),
- None,
- nulls,
- 0,
- vec![buffer],
- vec![],
- )
- };
- Ok(PrimitiveArray::<T>::from(data))
+ let values_buf = take_native(values.values(), indices);
+ let nulls = take_nulls(values.nulls(), indices);
+ Ok(PrimitiveArray::new(values_buf, nulls).with_data_type(values.data_type().clone()))
}
-fn take_bits<IndexType>(
- values: &Buffer,
- values_offset: usize,
- indices: &PrimitiveArray<IndexType>,
-) -> Result<Buffer, ArrowError>
-where
- IndexType: ArrowPrimitiveType,
- IndexType::Native: ToPrimitive,
-{
- let len = indices.len();
- let values_slice = values.as_slice();
- let mut output_buffer = MutableBuffer::new_null(len);
- let output_slice = output_buffer.as_slice_mut();
-
- let indices_has_nulls = indices.null_count() > 0;
+#[inline(never)]
Review Comment:
Likewise this may be worth a comment in the code as well
--
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
[GitHub] [arrow-rs] tustvold merged pull request #4405: Improve `take` kernel performance on primitive arrays, fix bad null index handling (#4404)
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold merged PR #4405:
URL: https://github.com/apache/arrow-rs/pull/4405
--
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
[GitHub] [arrow-rs] Dandandan commented on a diff in pull request #4405: Improve take primitive performance (#4404)
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on code in PR #4405:
URL: https://github.com/apache/arrow-rs/pull/4405#discussion_r1228115449
##########
arrow-select/src/take.rs:
##########
@@ -374,148 +243,91 @@ fn take_primitive<T, I>(
where
T: ArrowPrimitiveType,
I: ArrowPrimitiveType,
- I::Native: ToPrimitive,
{
- let indices_nulls = indices.nulls().filter(|x| x.null_count() > 0);
- let values_nulls = values.nulls().filter(|x| x.null_count() > 0);
-
- // note: this function should only panic when "an index is not null and out of bounds".
- // if the index is null, its value is undefined and therefore we should not read from it.
- let (buffer, nulls) = match (values_nulls, indices_nulls) {
- (None, None) => {
- // * no nulls
- // * all `indices.values()` are valid
- take_no_nulls(values.values(), indices.values())?
- }
- (Some(values_nulls), None) => {
- // * nulls come from `values` alone
- // * all `indices.values()` are valid
- take_values_nulls(values.values(), values_nulls, indices.values())?
- }
- (None, Some(indices_nulls)) => {
- // in this branch it is unsound to read and use `index.values()`,
- // as doing so is UB when they come from a null slot.
- take_indices_nulls(values.values(), indices.values(), indices_nulls)?
- }
- (Some(values_nulls), Some(indices_nulls)) => {
- // in this branch it is unsound to read and use `index.values()`,
- // as doing so is UB when they come from a null slot.
- take_values_indices_nulls(
- values.values(),
- values_nulls,
- indices.values(),
- indices_nulls,
- )?
- }
- };
-
- let data = unsafe {
- ArrayData::new_unchecked(
- values.data_type().clone(),
- indices.len(),
- None,
- nulls,
- 0,
- vec![buffer],
- vec![],
- )
- };
- Ok(PrimitiveArray::<T>::from(data))
+ let values_buf = take_native(values.values(), indices);
+ let nulls = take_nulls(values.nulls(), indices);
+ Ok(PrimitiveArray::new(values_buf, nulls).with_data_type(values.data_type().clone()))
}
-fn take_bits<IndexType>(
- values: &Buffer,
- values_offset: usize,
- indices: &PrimitiveArray<IndexType>,
-) -> Result<Buffer, ArrowError>
-where
- IndexType: ArrowPrimitiveType,
- IndexType::Native: ToPrimitive,
-{
- let len = indices.len();
- let values_slice = values.as_slice();
- let mut output_buffer = MutableBuffer::new_null(len);
- let output_slice = output_buffer.as_slice_mut();
-
- let indices_has_nulls = indices.null_count() > 0;
+#[inline(never)]
+fn take_nulls<I: ArrowPrimitiveType>(
+ values: Option<&NullBuffer>,
+ indices: &PrimitiveArray<I>,
+) -> Option<NullBuffer> {
+ match values.filter(|n| n.null_count() > 0) {
+ Some(n) => {
+ let buffer = take_bits(n.inner(), indices);
+ Some(NullBuffer::new(buffer)).filter(|n| n.null_count() > 0)
+ }
+ None => indices.nulls().cloned(),
+ }
+}
- if indices_has_nulls {
- indices
+#[inline(never)]
+fn take_native<T: ArrowNativeType, I: ArrowPrimitiveType>(
+ values: &[T],
+ indices: &PrimitiveArray<I>,
+) -> ScalarBuffer<T> {
+ match indices.nulls().filter(|n| n.null_count() > 0) {
+ Some(n) => indices
+ .values()
.iter()
.enumerate()
- .try_for_each::<_, Result<(), ArrowError>>(|(i, index)| {
- if let Some(index) = index {
- let index = ToPrimitive::to_usize(&index).ok_or_else(|| {
- ArrowError::ComputeError("Cast to usize failed".to_string())
- })?;
-
- if bit_util::get_bit(values_slice, values_offset + index) {
- bit_util::set_bit(output_slice, i);
- }
- }
-
- Ok(())
- })?;
- } else {
- indices
+ .map(|(idx, index)| match values.get(index.as_usize()) {
+ Some(v) => *v,
+ None => match n.is_null(idx) {
+ true => T::default(),
+ false => panic!("Out-of-bounds index {index:?}"),
+ },
+ })
+ .collect(),
+ None => indices
.values()
.iter()
- .enumerate()
- .try_for_each::<_, Result<(), ArrowError>>(|(i, index)| {
- let index = ToPrimitive::to_usize(index).ok_or_else(|| {
- ArrowError::ComputeError("Cast to usize failed".to_string())
- })?;
+ .map(|index| values[index.as_usize()])
+ .collect(),
+ }
+}
- if bit_util::get_bit(values_slice, values_offset + index) {
- bit_util::set_bit(output_slice, i);
- }
- Ok(())
- })?;
+#[inline(never)]
+fn take_bits<I: ArrowPrimitiveType>(
+ values: &BooleanBuffer,
+ indices: &PrimitiveArray<I>,
+) -> BooleanBuffer {
+ let len = indices.len();
+ let mut output_buffer = MutableBuffer::new_null(len);
+ let output_slice = output_buffer.as_slice_mut();
+
+ match indices.nulls().filter(|n| n.null_count() > 0) {
+ Some(nulls) => nulls.valid_indices().for_each(|idx| {
+ if values.value(indices.value(idx).as_usize()) {
+ bit_util::set_bit(output_slice, idx);
+ }
+ }),
+ None => indices.values().iter().enumerate().for_each(|(i, index)| {
+ if values.value(index.as_usize()) {
Review Comment:
Wasn't there a faster method to create a bitmap rather than doing `set_bit` in sequence?
--
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