jayzhan211 commented on code in PR #12623:
URL: https://github.com/apache/datafusion/pull/12623#discussion_r1776517883
##########
datafusion/physical-plan/src/aggregates/group_values/group_value_row.rs:
##########
@@ -121,37 +145,64 @@ impl<T: ArrowPrimitiveType> ArrayRowEq for
PrimitiveGroupValueBuilder<T> {
}
fn size(&self) -> usize {
- self.group_values.allocated_size() + self.nulls.allocated_size()
+ // BooleanBufferBuilder builder::capacity returns capacity in bits
(not bytes)
+ let null_builder_size = self
+ .nulls
+ .as_ref()
+ .map(|nulls| nulls.capacity() / 8)
+ .unwrap_or(0);
+ self.group_values.allocated_size() + null_builder_size
}
fn build(self: Box<Self>) -> ArrayRef {
- if self.has_null {
- Arc::new(PrimitiveArray::<T>::new(
- ScalarBuffer::from(self.group_values),
- Some(NullBuffer::from(self.nulls)),
- ))
- } else {
- Arc::new(PrimitiveArray::<T>::new(
- ScalarBuffer::from(self.group_values),
- None,
- ))
- }
+ let Self {
+ group_values,
+ nulls,
+ nullable: _,
+ } = *self;
+ let null_buffer = nulls.map(|mut nulls|
NullBuffer::from(nulls.finish()));
+
+ Arc::new(PrimitiveArray::<T>::new(
+ ScalarBuffer::from(group_values),
+ null_buffer,
+ ))
}
fn take_n(&mut self, n: usize) -> ArrayRef {
- if self.has_null {
- let first_n = self.group_values.drain(0..n).collect::<Vec<_>>();
- let first_n_nulls = self.nulls.drain(0..n).collect::<Vec<_>>();
- Arc::new(PrimitiveArray::<T>::new(
- ScalarBuffer::from(first_n),
- Some(NullBuffer::from(first_n_nulls)),
- ))
- } else {
- let first_n = self.group_values.drain(0..n).collect::<Vec<_>>();
- self.nulls.truncate(self.nulls.len() - n);
- Arc::new(PrimitiveArray::<T>::new(ScalarBuffer::from(first_n),
None))
- }
+ let first_n = self.group_values.drain(0..n).collect::<Vec<_>>();
+
+ let first_n_nulls = self.nulls.take().map(|nulls| {
+ let (first_n_nulls, remaining_nulls) = take_n_nulls(n, nulls);
+ self.nulls = Some(remaining_nulls);
+ first_n_nulls
+ });
+
+ Arc::new(PrimitiveArray::<T>::new(
+ ScalarBuffer::from(first_n),
+ first_n_nulls,
+ ))
+ }
+}
+
+/// Takes the first `n` nulls from the `nulls` buffer and
+///
+/// Returns
+/// * [`NullBuffer`] with the first `n` nulls
+/// * [`BooleanBufferBuilder`] with the remaining nulls
+fn take_n_nulls(
+ n: usize,
+ mut nulls: BooleanBufferBuilder,
+) -> (NullBuffer, BooleanBufferBuilder) {
+ // Copy over the values at n..len-1 values to the start of a new builder
Review Comment:
Another approach I'm thinking of is we only append the boolean and implement
`finish_slice()` to take arbitrary slice, so we can avoid *shifting bitmap*
which requires careful implementation
--
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]
