diff --git a/arrow-select/src/coalesce.rs b/arrow-select/src/coalesce.rs index 6f5378a8f5c8..d0638ea53526 100644 --- a/arrow-select/src/coalesce.rs +++ b/arrow-select/src/coalesce.rs @@ -31,18 +31,24 @@ use std::sync::Arc; // Originally From DataFusion's coalesce module: // https://github.com/apache/datafusion/blob/9d2f04996604e709ee440b65f41e7b882f50b788/datafusion/physical-plan/src/coalesce/mod.rs#L26-L25 +mod boolean; mod byte_view; mod fixed_size_binary; mod generic; mod primitive; +use boolean::InProgressBooleanArray; use byte_view::InProgressByteViewArray; use fixed_size_binary::InProgressFixedSizeBinaryArray; use generic::GenericInProgressArray; use primitive::InProgressPrimitiveArray; fn has_sparse_filter_copy(data_type: &DataType) -> bool { - data_type.is_primitive() || matches!(data_type, DataType::Utf8View | DataType::BinaryView) + data_type.is_primitive() + || matches!( + data_type, + DataType::Boolean | DataType::Utf8View | DataType::BinaryView + ) } /// Maximum selected row fraction for the fused sparse-filter copy path. @@ -783,6 +789,7 @@ fn create_in_progress_array(data_type: &DataType, batch_size: usize) -> Box { Box::new(InProgressByteViewArray::::new(batch_size)) } + DataType::Boolean => Box::new(InProgressBooleanArray::new(batch_size)), DataType::FixedSizeBinary(size) => { Box::new(InProgressFixedSizeBinaryArray::new(*size, batch_size)) } @@ -1771,7 +1778,7 @@ mod tests { true, )])); let coalescer = BatchCoalescer::new(boolean, 100); - assert!(coalescer.has_non_specialized_filter_columns); + assert!(!coalescer.has_non_specialized_filter_columns); } #[derive(Debug, Clone, PartialEq)] diff --git a/arrow-select/src/coalesce/boolean.rs b/arrow-select/src/coalesce/boolean.rs new file mode 100644 index 000000000000..17c76b2b6196 --- /dev/null +++ b/arrow-select/src/coalesce/boolean.rs @@ -0,0 +1,159 @@ +// 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 super::InProgressArray; +use crate::filter::FilterPredicate; +use arrow_array::cast::AsArray; +use arrow_array::{Array, ArrayRef, BooleanArray}; +use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, NullBufferBuilder}; +use arrow_schema::ArrowError; +use std::sync::Arc; + +#[derive(Debug)] +pub(crate) struct InProgressBooleanArray { + source: Option, + batch_size: usize, + values: Option, + nulls: NullBufferBuilder, +} + +impl InProgressBooleanArray { + pub(crate) fn new(batch_size: usize) -> Self { + Self { + source: None, + batch_size, + values: None, + nulls: NullBufferBuilder::new(batch_size), + } + } + + fn values(&mut self) -> &mut BooleanBufferBuilder { + self.values + .get_or_insert_with(|| BooleanBufferBuilder::new(self.batch_size)) + } + + fn append_array(&mut self, array: &BooleanArray) { + self.values().append_buffer(array.values()); + if let Some(nulls) = array.nulls() { + self.nulls.append_buffer(nulls); + } else { + self.nulls.append_n_non_nulls(array.len()); + } + } +} + +impl InProgressArray for InProgressBooleanArray { + fn set_source(&mut self, source: Option) { + self.source = source; + } + + fn copy_rows(&mut self, offset: usize, len: usize) -> Result<(), ArrowError> { + let source = self.source.as_ref().ok_or_else(|| { + ArrowError::InvalidArgumentError( + "Internal Error: InProgressBooleanArray: source not set".to_string(), + ) + })?; + let array = source.as_boolean().slice(offset, len); + self.append_array(&array); + Ok(()) + } + + fn copy_rows_by_filter_from( + &mut self, + source: ArrayRef, + filter: &FilterPredicate, + ) -> Result<(), ArrowError> { + let filtered = filter.filter(source.as_ref())?; + self.append_array(filtered.as_boolean()); + Ok(()) + } + + fn finish(&mut self) -> Result { + let values = self + .values + .take() + .map(|mut values| values.finish()) + .unwrap_or_else(|| BooleanBuffer::new_unset(0)); + let nulls = self.nulls.finish(); + self.nulls = NullBufferBuilder::new(self.batch_size); + Ok(Arc::new(BooleanArray::new(values, nulls))) + } + + fn size(&self) -> usize { + self.source + .as_ref() + .map_or(0, |source| source.get_array_memory_size()) + + self + .values + .as_ref() + .map_or(0, |values| values.capacity() / 8) + + self.nulls.allocated_size() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::filter::FilterBuilder; + + #[test] + fn copies_slices_and_sparse_filters() { + let source = Arc::new(BooleanArray::from(vec![ + Some(true), + None, + Some(false), + Some(true), + Some(false), + ])) as ArrayRef; + let mut in_progress = InProgressBooleanArray::new(8); + in_progress.set_source(Some(Arc::clone(&source))); + in_progress.copy_rows(1, 2).unwrap(); + in_progress.set_source(None); + + let filter = BooleanArray::from(vec![true, false, false, true, false]); + let predicate = FilterBuilder::new(&filter).build(); + in_progress + .copy_rows_by_filter_from(source, &predicate) + .unwrap(); + + let output = in_progress.finish().unwrap(); + let expected = BooleanArray::from(vec![None, Some(false), Some(true), Some(true)]); + assert_eq!(output.as_boolean(), &expected); + } + + #[test] + fn finish_resets_buffers() { + let mut in_progress = InProgressBooleanArray::new(8); + assert_eq!(in_progress.finish().unwrap().len(), 0); + + let source = Arc::new(BooleanArray::from(vec![ + Some(true), + None, + Some(false), + None, + ])); + in_progress.set_source(Some(Arc::clone(&source) as ArrayRef)); + in_progress.copy_rows(0, 2).unwrap(); + let first = in_progress.finish().unwrap(); + assert_eq!(first.as_boolean(), &source.slice(0, 2)); + + in_progress.copy_rows(2, 2).unwrap(); + let second = in_progress.finish().unwrap(); + assert_eq!(second.as_boolean(), &source.slice(2, 2)); + assert_eq!(in_progress.finish().unwrap().len(), 0); + } +}