kosiew commented on code in PR #17726: URL: https://github.com/apache/datafusion/pull/17726#discussion_r2383807325
########## datafusion/physical-plan/src/aggregates/group_values/multi_group_by/boolean.rs: ########## @@ -0,0 +1,475 @@ +// 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 std::sync::Arc; + +use crate::aggregates::group_values::multi_group_by::{nulls_equal_to, GroupColumn}; +use crate::aggregates::group_values::null_builder::MaybeNullBufferBuilder; +use arrow::array::{Array as _, ArrayRef, AsArray, BooleanArray, BooleanBufferBuilder}; +use datafusion_common::Result; +use itertools::izip; + +/// An implementation of [`GroupColumn`] for booleans +/// +/// Optimized to skip null buffer construction if the input is known to be non nullable +/// +/// # Template parameters +/// +/// `NULLABLE`: if the data can contain any nulls +#[derive(Debug)] +pub struct BooleanGroupValueBuilder<const NULLABLE: bool> { + buffer: BooleanBufferBuilder, + nulls: MaybeNullBufferBuilder, +} + +impl<const NULLABLE: bool> BooleanGroupValueBuilder<NULLABLE> { + /// Create a new `BooleanGroupValueBuilder` + pub fn new() -> Self { + Self { + buffer: BooleanBufferBuilder::new(0), + nulls: MaybeNullBufferBuilder::new(), + } + } +} + +impl<const NULLABLE: bool> GroupColumn for BooleanGroupValueBuilder<NULLABLE> { + fn equal_to(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool { + if NULLABLE { + let exist_null = self.nulls.is_null(lhs_row); + let input_null = array.is_null(rhs_row); + if let Some(result) = nulls_equal_to(exist_null, input_null) { + return result; + } + } + + self.buffer.get_bit(lhs_row) == array.as_boolean().value(rhs_row) + } + + fn append_val(&mut self, array: &ArrayRef, row: usize) -> Result<()> { + if NULLABLE { + if array.is_null(row) { + self.nulls.append(true); + self.buffer.append(bool::default()); + } else { + self.nulls.append(false); + self.buffer.append(array.as_boolean().value(row)); + } + } else { + self.buffer.append(array.as_boolean().value(row)); + } + + Ok(()) + } + + fn vectorized_equal_to( + &self, + lhs_rows: &[usize], + array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut [bool], + ) { + let array = array.as_boolean(); + + let iter = izip!( + lhs_rows.iter(), + rhs_rows.iter(), + equal_to_results.iter_mut(), + ); + + for (&lhs_row, &rhs_row, equal_to_result) in iter { + // Has found not equal to in previous column, don't need to check + if !*equal_to_result { + continue; + } + + if NULLABLE { + let exist_null = self.nulls.is_null(lhs_row); + let input_null = array.is_null(rhs_row); + if let Some(result) = nulls_equal_to(exist_null, input_null) { + *equal_to_result = result; + continue; + } + } + + *equal_to_result = self.buffer.get_bit(lhs_row) == array.value(rhs_row); + } + } + + fn vectorized_append(&mut self, array: &ArrayRef, rows: &[usize]) -> Result<()> { + let arr = array.as_boolean(); + + let null_count = array.null_count(); + let num_rows = array.len(); + let all_null_or_non_null = if null_count == 0 { + Some(true) + } else if null_count == num_rows { + Some(false) + } else { + None + }; + + match (NULLABLE, all_null_or_non_null) { + (true, None) => { + for &row in rows { + if array.is_null(row) { + self.nulls.append(true); + self.buffer.append(bool::default()); + } else { + self.nulls.append(false); + self.buffer.append(arr.value(row)); + } + } + } + + (true, Some(true)) => { + self.nulls.append_n(rows.len(), false); + for &row in rows { + self.buffer.append(arr.value(row)); + } + } + + (true, Some(false)) => { + self.nulls.append_n(rows.len(), true); + self.buffer.append_n(rows.len(), bool::default()); + } + + (false, _) => { + for &row in rows { + self.buffer.append(arr.value(row)); + } + } + } + + Ok(()) + } + + fn len(&self) -> usize { + self.buffer.len() + } + + fn size(&self) -> usize { + self.buffer.capacity() / 8 + self.nulls.allocated_size() + } + + fn build(self: Box<Self>) -> ArrayRef { + let Self { mut buffer, nulls } = *self; + + let nulls = nulls.build(); + if !NULLABLE { + assert!(nulls.is_none(), "unexpected nulls in non nullable input"); + } + + let arr = BooleanArray::new(buffer.finish(), nulls); + + Arc::new(arr) + } + + fn take_n(&mut self, n: usize) -> ArrayRef { + let first_n_nulls = if NULLABLE { self.nulls.take_n(n) } else { None }; + + let mut new_builder = BooleanBufferBuilder::new(self.buffer.len()); + new_builder.append_packed_range(n..self.buffer.len(), self.buffer.as_slice()); + std::mem::swap(&mut new_builder, &mut self.buffer); + + // take only first n values from the original builder + new_builder.truncate(n); + new_builder.finish(); + + Arc::new(BooleanArray::new(new_builder.finish(), first_n_nulls)) Review Comment: Why call new_builder.finish() twice? -- 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]
