alamb commented on code in PR #11627:
URL: https://github.com/apache/datafusion/pull/11627#discussion_r1697643230
##########
datafusion/functions-aggregate/src/count.rs:
##########
@@ -433,6 +433,49 @@ impl GroupsAccumulator for CountGroupsAccumulator {
Ok(vec![Arc::new(counts) as ArrayRef])
}
+ fn convert_to_state(
+ &self,
+ values: &[ArrayRef],
+ opt_filter: Option<&BooleanArray>,
+ ) -> Result<Vec<ArrayRef>> {
+ let values = &values[0];
+
+ let state_array = match (values.logical_nulls(), opt_filter) {
+ (Some(nulls), None) => {
+ let mut builder = Int64Builder::with_capacity(values.len());
+ nulls
+ .into_iter()
+ .for_each(|is_valid| builder.append_value(is_valid as
i64));
+ builder.finish()
+ }
+ (Some(nulls), Some(filter)) => {
+ let mut builder = Int64Builder::with_capacity(values.len());
+ nulls.into_iter().zip(filter.iter()).for_each(
+ |(is_valid, filter_value)| {
+ builder.append_value(
Review Comment:
I came up wtih this in https://github.com/apache/datafusion/pull/11734:
```rust
/// Converts a `BooleanBuffer` representing a filter to a `NullBuffer`
/// where the NullBuffer is true for all values that were true
/// in the filter and `null` for any values that were false or null
fn filter_to_nulls(filter: &BooleanArray) -> Option<NullBuffer> {
let (filter_bools, filter_nulls) = filter.clone().into_parts();
// Only keep values where the filter was true
// convert all false to null
let filter_bools = NullBuffer::from(filter_bools);
NullBuffer::union(Some(&filter_bools), filter_nulls.as_ref())
}
/// Compute the final null mask for an array
///
/// The output null mask :
/// * is true (non null) for all values that were true in the filter and non
null in the input
/// * is false (null) for all values that were false in the filter or null
in the input
fn filtered_null_mask(
opt_filter: Option<&BooleanArray>,
input: &dyn Array,
) -> Option<NullBuffer> {
let opt_filter = opt_filter.and_then(filter_to_nulls);
NullBuffer::union(opt_filter.as_ref(), input.nulls())
}
```
And then you compute the final null mask without messing with the input:
```rust
let nulls = filtered_null_mask(opt_filter, sums);
let sums = PrimitiveArray::<T>::new(sums.values().clone(), nulls)
.with_data_type(self.sum_data_type.clone());
```
--
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-unsubscr...@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org
