You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "JasonLi-cn (via GitHub)" <gi...@apache.org> on 2024/03/22 06:55:21 UTC
[PR] Optimization: concat function [arrow-datafusion]
JasonLi-cn opened a new pull request, #9732:
URL: https://github.com/apache/arrow-datafusion/pull/9732
## 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 #.
## Rationale for this change
optimize `concat` and `concat_ws` function.
### Benchmark(only concat)
```shell
Gnuplot not found, using plotters backend
concat function/concat(old)/1024
time: [91.400 µs 93.575 µs 96.505 µs]
Found 10 outliers among 100 measurements (10.00%)
5 (5.00%) high mild
5 (5.00%) high severe
concat function/concat(new)/1024
time: [18.160 µs 18.192 µs 18.230 µs]
Found 12 outliers among 100 measurements (12.00%)
1 (1.00%) low mild
4 (4.00%) high mild
7 (7.00%) high severe
concat function/concat(old)/4096
time: [88.002 µs 88.087 µs 88.185 µs]
Found 8 outliers among 100 measurements (8.00%)
1 (1.00%) low mild
7 (7.00%) high severe
concat function/concat(new)/4096
time: [18.303 µs 18.407 µs 18.538 µs]
Found 4 outliers among 100 measurements (4.00%)
3 (3.00%) high mild
1 (1.00%) high severe
concat function/concat(old)/8192
time: [88.308 µs 88.534 µs 88.798 µs]
Found 5 outliers among 100 measurements (5.00%)
2 (2.00%) high mild
3 (3.00%) high severe
concat function/concat(new)/8192
time: [18.229 µs 18.310 µs 18.413 µs]
Found 16 outliers among 100 measurements (16.00%)
5 (5.00%) high mild
11 (11.00%) high severe
```
### Attention
For the purpose of benchmarking, I haven't officially replaced the `concat` and `concat_ws` function yet. If the community finds this PR meaningful, I will proceed with the replacement.
## 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 these changes tested?
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example, are they covered by existing tests)?
-->
## 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 `api 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
Re: [PR] Optimization: concat function [arrow-datafusion]
Posted by "JasonLi-cn (via GitHub)" <gi...@apache.org>.
JasonLi-cn commented on code in PR #9732:
URL: https://github.com/apache/arrow-datafusion/pull/9732#discussion_r1545191216
##########
datafusion/physical-expr/src/string_expressions.rs:
##########
@@ -227,6 +229,132 @@ pub fn concat(args: &[ColumnarValue]) -> Result<ColumnarValue> {
}
}
+enum ColumnarValueRef<'a> {
+ Scalar(&'a [u8]),
+ Array(&'a StringArray),
+}
+
+impl<'a> ColumnarValueRef<'a> {
+ #[inline]
+ fn is_valid(&self, i: usize) -> bool {
+ match &self {
+ Self::Scalar(_) => true,
+ Self::Array(array) => array.is_valid(i),
+ }
+ }
+
+ #[inline]
+ fn nulls(&self) -> Option<NullBuffer> {
+ match &self {
+ Self::Scalar(_) => None,
+ Self::Array(array) => array.nulls().map(|b| b.clone()),
+ }
+ }
+}
+
+struct StringArrayBuilder {
+ offsets_buffer: MutableBuffer,
+ value_buffer: MutableBuffer,
+}
+
+impl StringArrayBuilder {
+ fn with_capacity(item_capacity: usize, data_capacity: usize) -> Self {
+ let mut offsets_buffer = MutableBuffer::with_capacity(
+ (item_capacity + 1) * std::mem::size_of::<i32>(),
+ );
+ unsafe { offsets_buffer.push_unchecked(0_i32) };
+ Self {
+ offsets_buffer,
+ value_buffer: MutableBuffer::with_capacity(data_capacity),
+ }
+ }
+
+ fn write<const CHECK_VALID: bool>(&mut self, column: &ColumnarValueRef, i: usize) {
+ match column {
+ ColumnarValueRef::Scalar(s) => {
+ self.value_buffer.extend_from_slice(s);
Review Comment:
1. My initial intention for optimizing this function was to avoid creating a new String and then using `push_str` each time in `concat` function, like:
```rust
let mut owned_string: String = "".to_owned();
for arg in args {
match arg {
ColumnarValue::Scalar(ScalarValue::Utf8(maybe_value)) => {
if let Some(value) = maybe_value {
owned_string.push_str(value);
}
}
ColumnarValue::Array(v) => {
if v.is_valid(index) {
let v = as_string_array(v).unwrap();
owned_string.push_str(v.value(index));
}
}
_ => unreachable!(),
}
}
Some(owned_string)
```
3. Additionally, by precalculating the expected length of the result, I avoided the need to reallocate memory.
4. I used the `extend_from_slice` function because I referred to the `append_slice` function of `BufferBuilder`.
```rust
#[inline]
pub fn append_slice(&mut self, slice: &[T]) {
self.buffer.extend_from_slice(slice);
self.len += slice.len();
}
```
--
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] Optimization: concat function [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #9732:
URL: https://github.com/apache/arrow-datafusion/pull/9732#discussion_r1535993449
##########
datafusion/physical-expr/src/string_expressions.rs:
##########
@@ -227,6 +229,132 @@ pub fn concat(args: &[ColumnarValue]) -> Result<ColumnarValue> {
}
}
+enum ColumnarValueRef<'a> {
+ Scalar(&'a [u8]),
+ Array(&'a StringArray),
+}
+
+impl<'a> ColumnarValueRef<'a> {
+ #[inline]
+ fn is_valid(&self, i: usize) -> bool {
+ match &self {
+ Self::Scalar(_) => true,
+ Self::Array(array) => array.is_valid(i),
+ }
+ }
+
+ #[inline]
+ fn nulls(&self) -> Option<NullBuffer> {
+ match &self {
+ Self::Scalar(_) => None,
+ Self::Array(array) => array.nulls().map(|b| b.clone()),
+ }
+ }
+}
+
+struct StringArrayBuilder {
Review Comment:
I think some comments that explained how this was different than https://docs.rs/arrow/latest/arrow/array/type.StringBuilder.html would help. Maybe simply a note that it didn't check UTF8 again?
I wonder if we could get the same effect by adding an `unsafe` function to `StringBuilder`, like
```rust
/// Adds bytes to the in progress string, without checking for valid utf8
///
/// Safety: requires that bytes are valid utf8, otherwise an invalid StringArray will result
unsafe fn append_unchecked(&mut self, bytes: &[u8])
```
And then using `StringBuilder` here
--
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] Optimization: concat function [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #9732:
URL: https://github.com/apache/arrow-datafusion/pull/9732#issuecomment-2037363093
Thanks again @JasonLi-cn 🙏
--
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] Optimization: concat function [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #9732:
URL: https://github.com/apache/arrow-datafusion/pull/9732#issuecomment-2015660661
I also filed https://github.com/apache/arrow-datafusion/issues/9742 to track this improvement
--
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] Optimization: concat function [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #9732:
URL: https://github.com/apache/arrow-datafusion/pull/9732#issuecomment-2016754241
Marking as draft as I think this PR is no longer waiting on feedback. Please mark it as ready for review when it is ready for another look
--
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] Optimization: concat function [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #9732:
URL: https://github.com/apache/arrow-datafusion/pull/9732#discussion_r1535996840
##########
datafusion/physical-expr/src/string_expressions.rs:
##########
@@ -227,6 +229,132 @@ pub fn concat(args: &[ColumnarValue]) -> Result<ColumnarValue> {
}
}
+enum ColumnarValueRef<'a> {
+ Scalar(&'a [u8]),
+ Array(&'a StringArray),
+}
+
+impl<'a> ColumnarValueRef<'a> {
+ #[inline]
+ fn is_valid(&self, i: usize) -> bool {
+ match &self {
+ Self::Scalar(_) => true,
+ Self::Array(array) => array.is_valid(i),
+ }
+ }
+
+ #[inline]
+ fn nulls(&self) -> Option<NullBuffer> {
+ match &self {
+ Self::Scalar(_) => None,
+ Self::Array(array) => array.nulls().map(|b| b.clone()),
+ }
+ }
+}
+
+struct StringArrayBuilder {
+ offsets_buffer: MutableBuffer,
+ value_buffer: MutableBuffer,
+}
+
+impl StringArrayBuilder {
+ fn with_capacity(item_capacity: usize, data_capacity: usize) -> Self {
+ let mut offsets_buffer = MutableBuffer::with_capacity(
+ (item_capacity + 1) * std::mem::size_of::<i32>(),
+ );
+ unsafe { offsets_buffer.push_unchecked(0_i32) };
+ Self {
+ offsets_buffer,
+ value_buffer: MutableBuffer::with_capacity(data_capacity),
+ }
+ }
+
+ fn write<const CHECK_VALID: bool>(&mut self, column: &ColumnarValueRef, i: usize) {
+ match column {
+ ColumnarValueRef::Scalar(s) => {
+ self.value_buffer.extend_from_slice(s);
Review Comment:
Is the primary speed savings gained from not checking UTF8 validity (and just copying byte slices)?
##########
datafusion/physical-expr/src/string_expressions.rs:
##########
@@ -227,6 +229,132 @@ pub fn concat(args: &[ColumnarValue]) -> Result<ColumnarValue> {
}
}
+enum ColumnarValueRef<'a> {
+ Scalar(&'a [u8]),
+ Array(&'a StringArray),
+}
+
+impl<'a> ColumnarValueRef<'a> {
+ #[inline]
+ fn is_valid(&self, i: usize) -> bool {
+ match &self {
+ Self::Scalar(_) => true,
+ Self::Array(array) => array.is_valid(i),
+ }
+ }
+
+ #[inline]
+ fn nulls(&self) -> Option<NullBuffer> {
+ match &self {
+ Self::Scalar(_) => None,
+ Self::Array(array) => array.nulls().map(|b| b.clone()),
+ }
+ }
+}
+
+struct StringArrayBuilder {
+ offsets_buffer: MutableBuffer,
+ value_buffer: MutableBuffer,
+}
+
+impl StringArrayBuilder {
+ fn with_capacity(item_capacity: usize, data_capacity: usize) -> Self {
+ let mut offsets_buffer = MutableBuffer::with_capacity(
+ (item_capacity + 1) * std::mem::size_of::<i32>(),
+ );
+ unsafe { offsets_buffer.push_unchecked(0_i32) };
+ Self {
+ offsets_buffer,
+ value_buffer: MutableBuffer::with_capacity(data_capacity),
+ }
+ }
+
+ fn write<const CHECK_VALID: bool>(&mut self, column: &ColumnarValueRef, i: usize) {
+ match column {
+ ColumnarValueRef::Scalar(s) => {
+ self.value_buffer.extend_from_slice(s);
+ }
+ ColumnarValueRef::Array(array) => {
+ if !CHECK_VALID || array.is_valid(i) {
+ self.value_buffer
+ .extend_from_slice(array.value(i).as_bytes());
+ }
+ }
+ }
+ }
+
+ fn append_offset(&mut self) {
+ let next_offset: i32 = self
+ .value_buffer
+ .len()
+ .try_into()
+ .expect("byte array offset overflow");
+ unsafe { self.offsets_buffer.push_unchecked(next_offset) };
+ }
+
+ fn finish(self, null_buffer: Option<NullBuffer>) -> StringArray {
+ let array_builder = ArrayDataBuilder::new(DataType::Utf8)
+ .len(self.offsets_buffer.len() / std::mem::size_of::<i32>() - 1)
+ .add_buffer(self.offsets_buffer.into())
+ .add_buffer(self.value_buffer.into())
+ .nulls(null_buffer);
+ let array_data = unsafe { array_builder.build_unchecked() };
+ StringArray::from(array_data)
+ }
+}
+
+pub fn concat2(args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ let array_len = args
+ .iter()
+ .filter_map(|x| match x {
+ ColumnarValue::Array(array) => Some(array.len()),
+ _ => None,
+ })
+ .next();
+
+ // Scalar
+ if array_len.is_none() {
+ let mut result = String::new();
+ for arg in args {
+ if let ColumnarValue::Scalar(ScalarValue::Utf8(Some(v))) = arg {
+ result.push_str(v);
+ }
+ }
+ return Ok(ColumnarValue::Scalar(ScalarValue::Utf8(Some(result))));
+ }
+
+ // Array
+ let len = array_len.unwrap();
+ let mut data_size = 0;
+ let mut columns = Vec::with_capacity(args.len());
+
+ for arg in args {
+ match arg {
+ ColumnarValue::Scalar(ScalarValue::Utf8(maybe_value)) => {
+ if let Some(s) = maybe_value {
+ data_size += s.len() * len;
+ columns.push(ColumnarValueRef::Scalar(s.as_bytes()));
+ }
+ }
+ ColumnarValue::Array(array) => {
+ let string_array = as_string_array(array)?;
+ data_size += string_array.values().len();
+ columns.push(ColumnarValueRef::Array(string_array));
+ }
+ _ => unreachable!(),
+ }
+ }
+
+ let mut builder = StringArrayBuilder::with_capacity(len, data_size);
+ for i in 0..len {
+ columns
+ .iter()
+ .for_each(|column| builder.write::<true>(column, i));
Review Comment:
You could also potentially special case when you know the null counts are 0 (aka there are no nulls in the column) and you can avoid checking `is_null`
##########
datafusion/physical-expr/src/string_expressions.rs:
##########
@@ -227,6 +229,132 @@ pub fn concat(args: &[ColumnarValue]) -> Result<ColumnarValue> {
}
}
+enum ColumnarValueRef<'a> {
+ Scalar(&'a [u8]),
+ Array(&'a StringArray),
+}
+
+impl<'a> ColumnarValueRef<'a> {
+ #[inline]
+ fn is_valid(&self, i: usize) -> bool {
+ match &self {
+ Self::Scalar(_) => true,
+ Self::Array(array) => array.is_valid(i),
+ }
+ }
+
+ #[inline]
+ fn nulls(&self) -> Option<NullBuffer> {
+ match &self {
+ Self::Scalar(_) => None,
+ Self::Array(array) => array.nulls().map(|b| b.clone()),
+ }
+ }
+}
+
+struct StringArrayBuilder {
Review Comment:
I think some comments that explained how this was different than https://docs.rs/arrow/latest/arrow/array/type.StringBuilder.html would help. Maybe simply a note that it didn't check UTF8 again?
I wonder if we could get the same effect by adding an `unsafe` function to `StringBuilder`, like
```rust
/// Adds bytes to the in progress string, without checking for valid utf8
///
/// Safety: requires that bytes are valid utf8, otherwise an invalid StringArray will result
unsafe fn append_unchecked(&mut self, bytes: &[u8])
##########
datafusion/physical-expr/benches/concat.rs:
##########
@@ -0,0 +1,50 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::util::bench_util::create_string_array_with_len;
+use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
+use datafusion_common::ScalarValue;
+use datafusion_expr::ColumnarValue;
+use datafusion_physical_expr::string_expressions::{concat, concat2};
+use std::sync::Arc;
+
+fn create_args(size: usize, str_len: usize) -> Vec<ColumnarValue> {
+ let array = Arc::new(create_string_array_with_len::<i32>(size, 0.2, str_len));
+ let scalar = ScalarValue::Utf8(Some(", ".to_string()));
+ vec![
+ ColumnarValue::Array(array.clone()),
+ ColumnarValue::Scalar(scalar),
+ ColumnarValue::Array(array),
+ ]
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+ for size in [1024, 4096, 8192] {
+ let args = create_args(1024, 32);
Review Comment:
did you mean to use `size` here too? Right now size seems only used for the name
```suggestion
let args = create_args(size, 32);
```
--
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] Optimization: concat function [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb merged PR #9732:
URL: https://github.com/apache/arrow-datafusion/pull/9732
--
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] Optimization: concat function [arrow-datafusion]
Posted by "JasonLi-cn (via GitHub)" <gi...@apache.org>.
JasonLi-cn commented on code in PR #9732:
URL: https://github.com/apache/arrow-datafusion/pull/9732#discussion_r1545195362
##########
datafusion/physical-expr/src/string_expressions.rs:
##########
@@ -227,6 +229,132 @@ pub fn concat(args: &[ColumnarValue]) -> Result<ColumnarValue> {
}
}
+enum ColumnarValueRef<'a> {
+ Scalar(&'a [u8]),
+ Array(&'a StringArray),
+}
+
+impl<'a> ColumnarValueRef<'a> {
+ #[inline]
+ fn is_valid(&self, i: usize) -> bool {
+ match &self {
+ Self::Scalar(_) => true,
+ Self::Array(array) => array.is_valid(i),
+ }
+ }
+
+ #[inline]
+ fn nulls(&self) -> Option<NullBuffer> {
+ match &self {
+ Self::Scalar(_) => None,
+ Self::Array(array) => array.nulls().map(|b| b.clone()),
+ }
+ }
+}
+
+struct StringArrayBuilder {
Review Comment:
Another issue is that we don't need `NullBufferBuilder` here, but `GenericByteBuilder` defaults to using `NullBufferBuilder`, which I believe introduces unnecessary overhead.
--
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] Optimization: concat function [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #9732:
URL: https://github.com/apache/arrow-datafusion/pull/9732#discussion_r1550469799
##########
datafusion/physical-expr/src/string_expressions.rs:
##########
@@ -38,75 +40,147 @@ use datafusion_common::{
};
use datafusion_expr::ColumnarValue;
+enum ColumnarValueRef<'a> {
+ Scalar(&'a [u8]),
+ NullableArray(&'a StringArray),
+ NonNullableArray(&'a StringArray),
+}
+
+impl<'a> ColumnarValueRef<'a> {
+ #[inline]
+ fn is_valid(&self, i: usize) -> bool {
+ match &self {
+ Self::Scalar(_) | Self::NonNullableArray(_) => true,
+ Self::NullableArray(array) => array.is_valid(i),
+ }
+ }
+
+ #[inline]
+ fn nulls(&self) -> Option<NullBuffer> {
+ match &self {
+ Self::Scalar(_) | Self::NonNullableArray(_) => None,
+ Self::NullableArray(array) => array.nulls().cloned(),
+ }
+ }
+}
+
+struct StringArrayBuilder {
+ offsets_buffer: MutableBuffer,
+ value_buffer: MutableBuffer,
+}
+
+impl StringArrayBuilder {
+ fn with_capacity(item_capacity: usize, data_capacity: usize) -> Self {
+ let mut offsets_buffer = MutableBuffer::with_capacity(
+ (item_capacity + 1) * std::mem::size_of::<i32>(),
+ );
+ unsafe { offsets_buffer.push_unchecked(0_i32) };
Review Comment:
Is it really necessary to avoid the bounds check for a single offset?
##########
datafusion/physical-expr/src/string_expressions.rs:
##########
@@ -38,75 +40,147 @@ use datafusion_common::{
};
use datafusion_expr::ColumnarValue;
+enum ColumnarValueRef<'a> {
+ Scalar(&'a [u8]),
+ NullableArray(&'a StringArray),
+ NonNullableArray(&'a StringArray),
+}
+
+impl<'a> ColumnarValueRef<'a> {
+ #[inline]
+ fn is_valid(&self, i: usize) -> bool {
+ match &self {
+ Self::Scalar(_) | Self::NonNullableArray(_) => true,
+ Self::NullableArray(array) => array.is_valid(i),
+ }
+ }
+
+ #[inline]
+ fn nulls(&self) -> Option<NullBuffer> {
+ match &self {
+ Self::Scalar(_) | Self::NonNullableArray(_) => None,
+ Self::NullableArray(array) => array.nulls().cloned(),
+ }
+ }
+}
+
+struct StringArrayBuilder {
Review Comment:
```suggestion
/// Optimized version of the StringBuilder in Arrow that:
/// 1. Precalculating the expected length of the result, avoiding reallocations.
/// 2. Avoids creating / incrementally creating a `NullBufferBuilder`
struct StringArrayBuilder {
```
##########
datafusion/physical-expr/src/string_expressions.rs:
##########
@@ -38,75 +40,147 @@ use datafusion_common::{
};
use datafusion_expr::ColumnarValue;
+enum ColumnarValueRef<'a> {
+ Scalar(&'a [u8]),
+ NullableArray(&'a StringArray),
+ NonNullableArray(&'a StringArray),
+}
+
+impl<'a> ColumnarValueRef<'a> {
+ #[inline]
+ fn is_valid(&self, i: usize) -> bool {
+ match &self {
+ Self::Scalar(_) | Self::NonNullableArray(_) => true,
+ Self::NullableArray(array) => array.is_valid(i),
+ }
+ }
+
+ #[inline]
+ fn nulls(&self) -> Option<NullBuffer> {
+ match &self {
+ Self::Scalar(_) | Self::NonNullableArray(_) => None,
+ Self::NullableArray(array) => array.nulls().cloned(),
+ }
+ }
+}
+
+struct StringArrayBuilder {
+ offsets_buffer: MutableBuffer,
+ value_buffer: MutableBuffer,
+}
+
+impl StringArrayBuilder {
+ fn with_capacity(item_capacity: usize, data_capacity: usize) -> Self {
+ let mut offsets_buffer = MutableBuffer::with_capacity(
+ (item_capacity + 1) * std::mem::size_of::<i32>(),
+ );
+ unsafe { offsets_buffer.push_unchecked(0_i32) };
+ Self {
+ offsets_buffer,
+ value_buffer: MutableBuffer::with_capacity(data_capacity),
+ }
+ }
+
+ fn write<const CHECK_VALID: bool>(&mut self, column: &ColumnarValueRef, i: usize) {
+ match column {
+ ColumnarValueRef::Scalar(s) => {
+ self.value_buffer.extend_from_slice(s);
+ }
+ ColumnarValueRef::NullableArray(array) => {
+ if !CHECK_VALID || array.is_valid(i) {
+ self.value_buffer
+ .extend_from_slice(array.value(i).as_bytes());
+ }
+ }
+ ColumnarValueRef::NonNullableArray(array) => {
+ self.value_buffer
+ .extend_from_slice(array.value(i).as_bytes());
+ }
+ }
+ }
+
+ fn append_offset(&mut self) {
+ let next_offset: i32 = self
+ .value_buffer
+ .len()
+ .try_into()
+ .expect("byte array offset overflow");
+ unsafe { self.offsets_buffer.push_unchecked(next_offset) };
+ }
+
+ fn finish(self, null_buffer: Option<NullBuffer>) -> StringArray {
+ let array_builder = ArrayDataBuilder::new(DataType::Utf8)
+ .len(self.offsets_buffer.len() / std::mem::size_of::<i32>() - 1)
+ .add_buffer(self.offsets_buffer.into())
+ .add_buffer(self.value_buffer.into())
+ .nulls(null_buffer);
+ let array_data = unsafe { array_builder.build_unchecked() };
Review Comment:
```suggestion
// SAFETY: all data that was appended was valid UTF8 and the values
// and offsets were created correctly
let array_data = unsafe { array_builder.build_unchecked() };
```
##########
datafusion/physical-expr/src/string_expressions.rs:
##########
@@ -38,75 +40,147 @@ use datafusion_common::{
};
use datafusion_expr::ColumnarValue;
+enum ColumnarValueRef<'a> {
+ Scalar(&'a [u8]),
+ NullableArray(&'a StringArray),
+ NonNullableArray(&'a StringArray),
+}
+
+impl<'a> ColumnarValueRef<'a> {
+ #[inline]
+ fn is_valid(&self, i: usize) -> bool {
+ match &self {
+ Self::Scalar(_) | Self::NonNullableArray(_) => true,
+ Self::NullableArray(array) => array.is_valid(i),
+ }
+ }
+
+ #[inline]
+ fn nulls(&self) -> Option<NullBuffer> {
+ match &self {
+ Self::Scalar(_) | Self::NonNullableArray(_) => None,
+ Self::NullableArray(array) => array.nulls().cloned(),
+ }
+ }
+}
+
+struct StringArrayBuilder {
+ offsets_buffer: MutableBuffer,
+ value_buffer: MutableBuffer,
+}
+
+impl StringArrayBuilder {
+ fn with_capacity(item_capacity: usize, data_capacity: usize) -> Self {
+ let mut offsets_buffer = MutableBuffer::with_capacity(
+ (item_capacity + 1) * std::mem::size_of::<i32>(),
+ );
+ unsafe { offsets_buffer.push_unchecked(0_i32) };
+ Self {
+ offsets_buffer,
+ value_buffer: MutableBuffer::with_capacity(data_capacity),
+ }
+ }
+
+ fn write<const CHECK_VALID: bool>(&mut self, column: &ColumnarValueRef, i: usize) {
+ match column {
+ ColumnarValueRef::Scalar(s) => {
+ self.value_buffer.extend_from_slice(s);
+ }
+ ColumnarValueRef::NullableArray(array) => {
+ if !CHECK_VALID || array.is_valid(i) {
+ self.value_buffer
+ .extend_from_slice(array.value(i).as_bytes());
+ }
+ }
+ ColumnarValueRef::NonNullableArray(array) => {
+ self.value_buffer
+ .extend_from_slice(array.value(i).as_bytes());
+ }
+ }
+ }
+
+ fn append_offset(&mut self) {
+ let next_offset: i32 = self
+ .value_buffer
+ .len()
+ .try_into()
+ .expect("byte array offset overflow");
+ unsafe { self.offsets_buffer.push_unchecked(next_offset) };
+ }
+
+ fn finish(self, null_buffer: Option<NullBuffer>) -> StringArray {
Review Comment:
I was trying to think if there was some way to create an invalid result with this API and I think the answer is No (even if append_offset wasn't called ever the offsets would still be valid.
--
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] Optimization: concat function [arrow-datafusion]
Posted by "JasonLi-cn (via GitHub)" <gi...@apache.org>.
JasonLi-cn commented on code in PR #9732:
URL: https://github.com/apache/arrow-datafusion/pull/9732#discussion_r1550883910
##########
datafusion/physical-expr/src/string_expressions.rs:
##########
@@ -38,75 +40,147 @@ use datafusion_common::{
};
use datafusion_expr::ColumnarValue;
+enum ColumnarValueRef<'a> {
+ Scalar(&'a [u8]),
+ NullableArray(&'a StringArray),
+ NonNullableArray(&'a StringArray),
+}
+
+impl<'a> ColumnarValueRef<'a> {
+ #[inline]
+ fn is_valid(&self, i: usize) -> bool {
+ match &self {
+ Self::Scalar(_) | Self::NonNullableArray(_) => true,
+ Self::NullableArray(array) => array.is_valid(i),
+ }
+ }
+
+ #[inline]
+ fn nulls(&self) -> Option<NullBuffer> {
+ match &self {
+ Self::Scalar(_) | Self::NonNullableArray(_) => None,
+ Self::NullableArray(array) => array.nulls().cloned(),
+ }
+ }
+}
+
+struct StringArrayBuilder {
+ offsets_buffer: MutableBuffer,
+ value_buffer: MutableBuffer,
+}
+
+impl StringArrayBuilder {
+ fn with_capacity(item_capacity: usize, data_capacity: usize) -> Self {
+ let mut offsets_buffer = MutableBuffer::with_capacity(
+ (item_capacity + 1) * std::mem::size_of::<i32>(),
+ );
+ unsafe { offsets_buffer.push_unchecked(0_i32) };
Review Comment:
Since it is safe here and there is a theoretical performance improvement, I have opted to use `push_unchecked` in this case.
--
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] Optimization: concat function [arrow-datafusion]
Posted by "JasonLi-cn (via GitHub)" <gi...@apache.org>.
JasonLi-cn commented on code in PR #9732:
URL: https://github.com/apache/arrow-datafusion/pull/9732#discussion_r1545186099
##########
datafusion/physical-expr/src/string_expressions.rs:
##########
@@ -227,6 +229,132 @@ pub fn concat(args: &[ColumnarValue]) -> Result<ColumnarValue> {
}
}
+enum ColumnarValueRef<'a> {
+ Scalar(&'a [u8]),
+ Array(&'a StringArray),
+}
+
+impl<'a> ColumnarValueRef<'a> {
+ #[inline]
+ fn is_valid(&self, i: usize) -> bool {
+ match &self {
+ Self::Scalar(_) => true,
+ Self::Array(array) => array.is_valid(i),
+ }
+ }
+
+ #[inline]
+ fn nulls(&self) -> Option<NullBuffer> {
+ match &self {
+ Self::Scalar(_) => None,
+ Self::Array(array) => array.nulls().map(|b| b.clone()),
+ }
+ }
+}
+
+struct StringArrayBuilder {
Review Comment:
The `Write` trait impl of StringBuilder can meet my current needs, but it is not very convenient to use, so I've defined a StringArrayBuilder. I agree with your suggestion to add an unsafe function to StringBuilder.
https://github.com/apache/arrow-rs/blob/9f36c883459405ecd9a5f4fdfa9a3317ab52302c/arrow-array/src/builder/generic_bytes_builder.rs#L231-L257
```rust
let mut builder = GenericStringBuilder::<i32>::new();
// Write data
write!(builder, "foo").unwrap();
write!(builder, "bar").unwrap();
// Finish value
builder.append_value("baz");
// Write second value
write!(builder, "v2").unwrap();
builder.append_value("");
let array = builder.finish();
assert_eq!(array.value(0), "foobarbaz");
assert_eq!(array.value(1), "v2");
```
--
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] Optimization: concat function [arrow-datafusion]
Posted by "JasonLi-cn (via GitHub)" <gi...@apache.org>.
JasonLi-cn commented on code in PR #9732:
URL: https://github.com/apache/arrow-datafusion/pull/9732#discussion_r1545191216
##########
datafusion/physical-expr/src/string_expressions.rs:
##########
@@ -227,6 +229,132 @@ pub fn concat(args: &[ColumnarValue]) -> Result<ColumnarValue> {
}
}
+enum ColumnarValueRef<'a> {
+ Scalar(&'a [u8]),
+ Array(&'a StringArray),
+}
+
+impl<'a> ColumnarValueRef<'a> {
+ #[inline]
+ fn is_valid(&self, i: usize) -> bool {
+ match &self {
+ Self::Scalar(_) => true,
+ Self::Array(array) => array.is_valid(i),
+ }
+ }
+
+ #[inline]
+ fn nulls(&self) -> Option<NullBuffer> {
+ match &self {
+ Self::Scalar(_) => None,
+ Self::Array(array) => array.nulls().map(|b| b.clone()),
+ }
+ }
+}
+
+struct StringArrayBuilder {
+ offsets_buffer: MutableBuffer,
+ value_buffer: MutableBuffer,
+}
+
+impl StringArrayBuilder {
+ fn with_capacity(item_capacity: usize, data_capacity: usize) -> Self {
+ let mut offsets_buffer = MutableBuffer::with_capacity(
+ (item_capacity + 1) * std::mem::size_of::<i32>(),
+ );
+ unsafe { offsets_buffer.push_unchecked(0_i32) };
+ Self {
+ offsets_buffer,
+ value_buffer: MutableBuffer::with_capacity(data_capacity),
+ }
+ }
+
+ fn write<const CHECK_VALID: bool>(&mut self, column: &ColumnarValueRef, i: usize) {
+ match column {
+ ColumnarValueRef::Scalar(s) => {
+ self.value_buffer.extend_from_slice(s);
Review Comment:
1. My initial intention for optimizing this function was to avoid creating a new String and then using `push_str` each time in `concat` function, like:
```rust
let mut owned_string: String = "".to_owned();
for arg in args {
match arg {
ColumnarValue::Scalar(ScalarValue::Utf8(maybe_value)) => {
if let Some(value) = maybe_value {
owned_string.push_str(value);
}
}
ColumnarValue::Array(v) => {
if v.is_valid(index) {
let v = as_string_array(v).unwrap();
owned_string.push_str(v.value(index));
}
}
_ => unreachable!(),
}
}
Some(owned_string)
```
3. Additionally, by precalculating the expected length of the result, I avoided the need to reallocate memory.
4. I used the `extend_from_slice` function because I referred to the `append_slice` function of `BufferBuilder`.
```rust
#[inline]
pub fn append_slice(&mut self, slice: &[T]) {
self.buffer.extend_from_slice(slice);
self.len += slice.len();
}
```
--
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