comphead commented on code in PR #19996:
URL: https://github.com/apache/datafusion/pull/19996#discussion_r2725949957
##########
datafusion/functions-nested/src/remove.rs:
##########
@@ -377,73 +377,83 @@ fn general_remove<OffsetSize: OffsetSizeTrait>(
);
}
};
- let data_type = list_field.data_type();
- let mut new_values = vec![];
+ let original_data = list_array.values().to_data();
// Build up the offsets for the final output array
let mut offsets = Vec::<OffsetSize>::with_capacity(arr_n.len() + 1);
offsets.push(OffsetSize::zero());
- // n is the number of elements to remove in this row
- for (row_index, (list_array_row, n)) in
- list_array.iter().zip(arr_n.iter()).enumerate()
- {
- match list_array_row {
- Some(list_array_row) => {
- let eq_array = utils::compare_element_to_list(
- &list_array_row,
- element_array,
- row_index,
- false,
- )?;
-
- // We need to keep at most first n elements as `false`, which
represent the elements to remove.
- let eq_array = if eq_array.false_count() < *n as usize {
- eq_array
- } else {
- let mut count = 0;
- eq_array
- .iter()
- .map(|e| {
- // Keep first n `false` elements, and reverse
other elements to `true`.
- if let Some(false) = e {
- if count < *n {
- count += 1;
- e
- } else {
- Some(true)
- }
- } else {
- e
- }
- })
- .collect::<BooleanArray>()
- };
-
- let filtered_array = arrow::compute::filter(&list_array_row,
&eq_array)?;
- offsets.push(
- offsets[row_index] +
OffsetSize::usize_as(filtered_array.len()),
- );
- new_values.push(filtered_array);
- }
- None => {
- // Null element results in a null row (no new offsets)
- offsets.push(offsets[row_index]);
+ let mut mutable = MutableArrayData::with_capacities(
+ vec![&original_data],
+ false,
+ Capacities::Array(original_data.len()),
+ );
+ let mut valid = NullBufferBuilder::new(list_array.len());
+
+ for (row_index, offset_window) in
list_array.offsets().windows(2).enumerate() {
+ if list_array.is_null(row_index) {
+ offsets.push(offsets[row_index]);
+ valid.append_null();
+ continue;
+ }
+
+ let start = offset_window[0].to_usize().unwrap();
+ let end = offset_window[1].to_usize().unwrap();
+ // n is the number of elements to remove in this row
+ let n = arr_n[row_index];
+
+ // compare each element in the list, `false` means the element matches
and should be removed
+ let eq_array = utils::compare_element_to_list(
+ &list_array.value(row_index),
+ element_array,
+ row_index,
+ false,
+ )?;
+
+ let false_count = eq_array.false_count();
Review Comment:
```suggestion
let num_to_remove = eq_array.false_count();
```
WDYT?
##########
datafusion/functions-nested/benches/array_remove.rs:
##########
@@ -0,0 +1,572 @@
+// 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.
+
+#[macro_use]
+extern crate criterion;
+
+use arrow::array::{
+ Array, ArrayRef, BinaryArray, BooleanArray, Decimal128Array,
FixedSizeBinaryArray,
+ Float64Array, Int64Array, ListArray, StringArray,
+};
+use arrow::buffer::OffsetBuffer;
+use arrow::datatypes::{DataType, Field};
+use criterion::{BenchmarkId, Criterion};
+use datafusion_common::ScalarValue;
+use datafusion_common::config::ConfigOptions;
+use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
+use datafusion_functions_nested::remove::ArrayRemove;
+use rand::Rng;
+use rand::SeedableRng;
+use rand::rngs::StdRng;
+use std::hint::black_box;
+use std::sync::Arc;
+
+const NUM_ROWS: usize = 10000;
+const ARRAY_SIZES: &[usize] = &[10, 100, 500];
+const SEED: u64 = 42;
+const NULL_DENSITY: f64 = 0.1;
+
+fn criterion_benchmark(c: &mut Criterion) {
+ // Test array_remove with different data types and array sizes
+ bench_array_remove_int64(c);
Review Comment:
great, would be also nice to see performance for nested datatypes, but it
can be done in the following PR, for now please mention a TODO here
##########
datafusion/functions-nested/src/remove.rs:
##########
@@ -377,73 +377,83 @@ fn general_remove<OffsetSize: OffsetSizeTrait>(
);
}
};
- let data_type = list_field.data_type();
- let mut new_values = vec![];
+ let original_data = list_array.values().to_data();
// Build up the offsets for the final output array
let mut offsets = Vec::<OffsetSize>::with_capacity(arr_n.len() + 1);
offsets.push(OffsetSize::zero());
- // n is the number of elements to remove in this row
- for (row_index, (list_array_row, n)) in
- list_array.iter().zip(arr_n.iter()).enumerate()
- {
- match list_array_row {
- Some(list_array_row) => {
- let eq_array = utils::compare_element_to_list(
- &list_array_row,
- element_array,
- row_index,
- false,
- )?;
-
- // We need to keep at most first n elements as `false`, which
represent the elements to remove.
- let eq_array = if eq_array.false_count() < *n as usize {
- eq_array
- } else {
- let mut count = 0;
- eq_array
- .iter()
- .map(|e| {
- // Keep first n `false` elements, and reverse
other elements to `true`.
- if let Some(false) = e {
- if count < *n {
- count += 1;
- e
- } else {
- Some(true)
- }
- } else {
- e
- }
- })
- .collect::<BooleanArray>()
- };
-
- let filtered_array = arrow::compute::filter(&list_array_row,
&eq_array)?;
- offsets.push(
- offsets[row_index] +
OffsetSize::usize_as(filtered_array.len()),
- );
- new_values.push(filtered_array);
- }
- None => {
- // Null element results in a null row (no new offsets)
- offsets.push(offsets[row_index]);
+ let mut mutable = MutableArrayData::with_capacities(
+ vec![&original_data],
+ false,
+ Capacities::Array(original_data.len()),
+ );
+ let mut valid = NullBufferBuilder::new(list_array.len());
+
+ for (row_index, offset_window) in
list_array.offsets().windows(2).enumerate() {
+ if list_array.is_null(row_index) {
+ offsets.push(offsets[row_index]);
+ valid.append_null();
+ continue;
+ }
+
+ let start = offset_window[0].to_usize().unwrap();
+ let end = offset_window[1].to_usize().unwrap();
+ // n is the number of elements to remove in this row
+ let n = arr_n[row_index];
+
+ // compare each element in the list, `false` means the element matches
and should be removed
+ let eq_array = utils::compare_element_to_list(
+ &list_array.value(row_index),
+ element_array,
+ row_index,
+ false,
+ )?;
+
+ let false_count = eq_array.false_count();
+
+ // Fast path: no elements to remove, copy entire row
+ if false_count == 0 {
+ mutable.extend(0, start, end);
+ offsets.push(offsets[row_index] + OffsetSize::usize_as(end -
start));
+ valid.append_non_null();
+ continue;
+ }
+
+ // Remove at most `n` matching elements
+ let max_removals = n.min(false_count as i64);
+ let mut removed = 0i64;
+ let mut copied = 0usize;
+ let mut batch_start: Option<usize> = None;
Review Comment:
```suggestion
// marks the beginning of a range of elements pending to be
copied.
let mut pending_batch_to_retain: Option<usize> = None;
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
