perf(arrow): specialize Boolean filtering

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<dyn
         DataType::BinaryView => {
             Box::new(InProgressByteViewArray::<BinaryViewType>::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)]

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<ArrayRef>,
+    batch_size: usize,
+    values: Option<BooleanBufferBuilder>,
+    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<ArrayRef>) {
+        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<ArrayRef, ArrowError> {
+        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);
+    }
+}