alamb commented on code in PR #12623:
URL: https://github.com/apache/datafusion/pull/12623#discussion_r1777745457
##########
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:
> The challenge is to "shift bitmap" without setting bit one by one but
byte-wise. I think we need get_bit_slice that can read the partial byte or full
byte efficiently and use set_bit_slice to copy those bytes to a new one.
> so we can avoid shifting bitmap which requires careful implementation.
There is already a function that does this shifting in arrow-rs (that also
has had quite a bit of careful optimizaton and testing already):
https://docs.rs/arrow/latest/arrow/util/bit_mask/fn.set_bits.html
To use it though you need mutable access to the underlying `[u8]`
I'll see what I can come up with tomorro
--
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
