jayzhan211 commented on code in PR #25707:
URL: https://github.com/apache/datafusion/pull/25707#discussion_r4099998730
##########
datafusion/physical-expr/src/aggregate.rs:
##########
@@ -670,6 +686,7 @@ pub struct AggregateFunctionExpr {
is_reversed: bool,
input_fields: Vec<FieldRef>,
is_nullable: bool,
+ batch_size: Option<usize>,
Review Comment:
The block size shouldn't be set at plan time on `AggregateFunctionExpr`. The
table uses the runtime `context.session_config().batch_size()`, so a plan
executed with a different `TaskContext` fails at the assert in
`aggregate_hash_table/common.rs:100` with `Internal("... left: 4096, right:
8192: Block size mismatch ...")`. Suggest dropping the field and the builder
method, and having the caller pass the runtime block size in:
```rs
pub fn blocked_groups_accumulator_supported(&self, block_size: usize) ->
bool {
```
```rs
pub fn create_blocked_groups_accumulator(
&self,
block_size: usize,
) -> Result<Box<dyn BlockedGroupsAccumulator>> {
```
##########
datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs:
##########
@@ -80,6 +84,38 @@ pub(in crate::aggregates) fn create_group_accumulator(
}
}
+/// Create an accumulator for `agg_expr` -- a [`BlockedGroupsAccumulator`] if
+/// that is supported by the aggregate, or a
+/// [`BlockedGroupsAccumulatorAdapter`] if not.
+pub(in crate::aggregates) fn create_blocked_group_accumulator(
Review Comment:
Should the blocked API be opt-in rather than the route for every
aggregation? A flat `GroupsAccumulator`/`GroupValues` can't serve
`BlockedEmitTo` cheaply: the adapters loop `EmitTo::First(block_size)`, and
each call shifts the whole table. Suggest taking the blocked path only when the
group values and every accumulator support it natively, and keeping the flat
path unchanged otherwise. That also keeps the adapters out of the public shape.
##########
datafusion/expr-common/src/blocked_groups_accumulator.rs:
##########
@@ -0,0 +1,398 @@
+// 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.
+
+//! Vectorized [`GroupsAccumulator`]
+
+use crate::groups_accumulator::EmitTo;
+use arrow::array::{ArrayRef, BooleanArray};
+use datafusion_common::{
+ Result, assert_ne_or_internal_err, assert_or_internal_err, exec_err,
not_impl_err,
+};
+use std::cmp::Ordering;
+
+/// Selects groups for a non-destructive grouped aggregation read.
+///
+/// Unlike [`BlockedEmitTo`], this selection does not remove groups or change
their
+/// indices. Selections created by [`Self::try_from_indices`] preserve the
+/// requested order and support duplicate indices.
+///
+/// A selection is validated once when it is constructed and can then be reused
+/// for the group values and accumulators participating in the same snapshot.
+/// Construct a new selection if the number or indexing of those groups
changes.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct BlockedGroupSelection<'a> {
+ block_size: usize,
+ total_num_groups: usize,
+ indices: Option<&'a [BlocksIndex]>,
+}
+
+impl<'a> BlockedGroupSelection<'a> {
+ /// Selects all `total_num_groups` groups in group-index order.
+ pub fn all(total_num_groups: usize, block_size: usize) -> Self {
+ Self {
+ total_num_groups,
+ indices: None,
+ block_size,
+ }
+ }
+
+ /// Selects groups in the order specified by `indices`.
+ ///
+ /// Returns an error if an index is not less than `total_num_groups`. Empty
+ /// selections are valid, and duplicate indices are preserved.
+ pub fn try_from_indices(
+ indices: &'a [BlocksIndex],
+ total_num_groups: usize,
+ block_size: usize,
+ ) -> Result<Self> {
+ let total_num_groups_parsed =
+ BlocksIndex::from_index_in_fixed_block_size(total_num_groups,
block_size);
+ if let Some(index) = indices
+ .iter()
+ .find(|&&index| index >= total_num_groups_parsed)
+ {
+ return exec_err!(
+ "Group index {index:?} is out of bounds for
{total_num_groups_parsed:?} ({total_num_groups}) groups"
+ );
+ }
+ Ok(Self {
+ total_num_groups,
+ indices: Some(indices),
+ block_size,
+ })
+ }
+
+ /// Returns the group count against which this selection was constructed.
+ pub fn total_num_groups(self) -> usize {
+ self.total_num_groups
+ }
+
+ /// Ensures this selection is being applied to the same number of groups
+ /// against which it was constructed.
+ ///
+ /// Preserving-read implementations should call this method with their
+ /// stored group count before using [`Self::iter`]. This check is `O(1)`;
+ /// the selected indices were already checked by
[`Self::try_from_indices`].
+ pub fn validate_num_groups(self, actual_num_groups: usize) -> Result<()> {
+ if actual_num_groups != self.total_num_groups {
+ return exec_err!(
+ "Group selection was constructed for {} groups but applied to
{actual_num_groups} groups",
+ self.total_num_groups
+ );
+ }
+ Ok(())
+ }
+
+ /// Returns the number of selected groups.
+ pub fn len(self) -> usize {
+ self.indices
+ .map_or(self.total_num_groups, |indices| indices.len())
+ }
+
+ /// Returns `true` if no groups are selected.
+ pub fn is_empty(self) -> bool {
+ self.len() == 0
+ }
+
+ /// Returns the selected group indices in output order.
+ pub fn iter(self) -> impl Iterator<Item = BlocksIndex> + 'a {
+ let (all, indices): (_, &'a [BlocksIndex]) = match self.indices {
+ None => (0..self.total_num_groups, &[]),
+ Some(indices) => (0..0, indices),
+ };
+ let block_size = self.block_size;
+ all.map(move |index| {
+ BlocksIndex::from_index_in_fixed_block_size(index, block_size)
+ })
+ .chain(indices.iter().copied())
+ }
+
+ #[doc(hidden)]
+ pub fn indices(&self) -> Option<&[BlocksIndex]> {
+ self.indices
+ }
+}
+
+/// Two `u32`s (8 bytes) rather than two `usize`s: an index is stored per
group in every blocked
+/// hash table entry, so its size decides the size of the tables. Both halves
fit `u32` by a wide
+/// margin, `index_in_block < block_size` (the batch size) and the number of
blocks is
+/// `groups / block_size`. The API keeps taking and returning `usize`.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Hash)]
+pub struct BlocksIndex {
+ block_index: u32,
+ index_in_block: u32,
+}
+
+impl BlocksIndex {
+ pub const ZERO: Self = Self {
+ // flat_index: 0,
+ block_index: 0,
+ index_in_block: 0,
+ };
+ pub const MAX: Self = Self {
+ // flat_index: usize::MAX,
+ block_index: u32::MAX,
+ index_in_block: u32::MAX,
+ };
+
+ pub fn new(block_index: usize, index_in_block: usize) -> Self {
+ Self {
+ block_index: Self::narrow(block_index),
+ index_in_block: Self::narrow(index_in_block),
+ }
+ }
+
+ #[inline(always)]
+ fn narrow(value: usize) -> u32 {
+ debug_assert!(
+ u32::try_from(value).is_ok(),
+ "block index part {value} does not fit in u32"
+ );
+ value as u32
+ }
+
+ #[inline(always)]
+ pub fn new_in_first_block(index_in_block: usize) -> Self {
+ // Implementation note:
+ // not having From<usize> that will do this instead even when it will
be more convenient
+ // so we can later change the layout to be a single usize with bit
shifts
+ Self::new(0, index_in_block)
+ // Self {
+ // flat_index: index_in_block,
+ // }
+ }
+
+ #[inline(always)]
+ pub fn from_index_in_fixed_block_size(index: usize, block_size: usize) ->
Self {
+ Self::new(index / block_size, index % block_size)
+ }
+
+ #[inline(always)]
+ pub fn into_index_in_fixed_block_size(self, block_size: usize) -> usize {
+ self.block_index as usize * block_size + self.index_in_block as usize
+ }
+
+ #[inline(always)]
+ pub fn block_index(&self) -> usize {
+ self.block_index as usize
+ }
+
+ #[inline(always)]
+ pub fn index_in_block(&self) -> usize {
+ self.index_in_block as usize
+ }
+
+ pub fn sub_flat_checked(self, rhs_flat: usize, block_size: usize) ->
Option<Self> {
+ let self_flat = self.into_index_in_fixed_block_size(block_size);
+ self_flat
+ .checked_sub(rhs_flat)
+ .map(|v| Self::from_index_in_fixed_block_size(v, block_size))
+ }
+
+ pub fn sub_flat(self, rhs_flat: usize, block_size: usize) -> Self {
+ Self::from_index_in_fixed_block_size(
+ self.into_index_in_fixed_block_size(block_size) - rhs_flat,
+ block_size,
+ )
+ }
+
+ pub fn gte_flat(self, rhs_flat: usize, block_size: usize) -> bool {
+ self.into_index_in_fixed_block_size(block_size) >= rhs_flat
+ }
+
+ pub fn prev_block(mut self) -> Self {
+ self.block_index -= 1;
+
+ self
+ }
+
+ pub fn prev_block_checked(mut self) -> Option<Self> {
+ self.block_index.checked_sub(1).map(|v| {
+ self.block_index = v;
+ self
+ })
+ }
+
+ pub fn add(self, n: usize, block_size: usize) -> BlocksIndex {
+ if self.index_in_block + n as u32 >= block_size as u32 {
+ let new_flat = self.into_index_in_fixed_block_size(block_size) + n;
+ Self::from_index_in_fixed_block_size(new_flat, block_size)
+ } else {
+ Self::new(self.block_index as usize, self.index_in_block as usize
+ n)
+ }
+ }
+
+ pub fn increment(mut self, block_size: usize) -> Self {
+ if self.index_in_block + 1 >= block_size as u32 {
+ self.block_index += 1;
+ self.index_in_block = 0;
+ } else {
+ self.index_in_block += 1;
+ }
+
+ self
+ }
+}
+
+impl PartialOrd for BlocksIndex {
+ fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+ Some(self.cmp(other))
+ }
+}
+
+impl Ord for BlocksIndex {
+ fn cmp(&self, other: &Self) -> Ordering {
+ self.block_index
+ .cmp(&other.block_index)
+ .then(self.index_in_block.cmp(&other.index_in_block))
+ }
+}
+
+/// Describes how many rows should be emitted during grouping.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum BlockedEmitTo {
+ All,
+ /// Emit next group
+ NextBlock,
+ /// Emit only the first `n` groups and shift all existing group
+ /// indexes down by `n`.
+ ///
+ /// For example, if `n=10`, group_index `0, 1, ... 9` are emitted
+ /// and group indexes `10, 11, 12, ...` become `0, 1, 2, ...`.
+ ///
+ /// Requirements:
+ /// 1. `n` must be smaller than block_size
+ /// 2. `n` is not 0
+ First(usize),
+}
+
+impl BlockedEmitTo {
+ pub fn into_emit_to(
+ self,
+ len: usize,
+ block_size: usize,
+ ) -> Result<impl Iterator<Item = EmitTo> + Clone> {
+ let mut to_emit = match self {
+ Self::All => len,
+ Self::NextBlock => len.min(block_size),
+ Self::First(n) => {
+ assert_ne_or_internal_err!(n, 0);
+ assert_or_internal_err!(
+ n <= len,
+ "n ({n}) must be less than or equal current length ({})",
+ len
+ );
+ assert_or_internal_err!(
+ n < block_size,
+ "n ({n}) must be less than current block size ({})",
+ block_size
+ );
+ n
+ }
+ };
+
+ let mut blocks = vec![];
+
+ let emitting_all = len == to_emit;
+
+ while to_emit > block_size {
+ blocks.push(EmitTo::First(block_size));
+ to_emit -= block_size
+ }
+
+ if to_emit > 0 && emitting_all {
+ blocks.push(EmitTo::All);
+ } else if to_emit > 0 {
+ blocks.push(EmitTo::First(to_emit));
+ }
+
+ Ok(blocks.into_iter())
+ }
+}
+
+pub trait BlockedGroupsAccumulator: Send + std::any::Any {
Review Comment:
This trait repeats `GroupsAccumulator` method for method, and
`BlockedGroupSelection` (L37) repeats `GroupSelection`; only the index type
changes. Could the two share a generic index type, or could the blocked methods
be defaulted on `GroupsAccumulator`? That would avoid two parallel public
traits that have to stay in sync.
##########
datafusion/expr-common/src/blocked_groups_accumulator.rs:
##########
@@ -0,0 +1,398 @@
+// 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.
+
+//! Vectorized [`GroupsAccumulator`]
+
+use crate::groups_accumulator::EmitTo;
+use arrow::array::{ArrayRef, BooleanArray};
+use datafusion_common::{
+ Result, assert_ne_or_internal_err, assert_or_internal_err, exec_err,
not_impl_err,
+};
+use std::cmp::Ordering;
+
+/// Selects groups for a non-destructive grouped aggregation read.
+///
+/// Unlike [`BlockedEmitTo`], this selection does not remove groups or change
their
+/// indices. Selections created by [`Self::try_from_indices`] preserve the
+/// requested order and support duplicate indices.
+///
+/// A selection is validated once when it is constructed and can then be reused
+/// for the group values and accumulators participating in the same snapshot.
+/// Construct a new selection if the number or indexing of those groups
changes.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct BlockedGroupSelection<'a> {
+ block_size: usize,
+ total_num_groups: usize,
+ indices: Option<&'a [BlocksIndex]>,
+}
+
+impl<'a> BlockedGroupSelection<'a> {
+ /// Selects all `total_num_groups` groups in group-index order.
+ pub fn all(total_num_groups: usize, block_size: usize) -> Self {
+ Self {
+ total_num_groups,
+ indices: None,
+ block_size,
+ }
+ }
+
+ /// Selects groups in the order specified by `indices`.
+ ///
+ /// Returns an error if an index is not less than `total_num_groups`. Empty
+ /// selections are valid, and duplicate indices are preserved.
+ pub fn try_from_indices(
+ indices: &'a [BlocksIndex],
+ total_num_groups: usize,
+ block_size: usize,
+ ) -> Result<Self> {
+ let total_num_groups_parsed =
+ BlocksIndex::from_index_in_fixed_block_size(total_num_groups,
block_size);
+ if let Some(index) = indices
+ .iter()
+ .find(|&&index| index >= total_num_groups_parsed)
+ {
+ return exec_err!(
+ "Group index {index:?} is out of bounds for
{total_num_groups_parsed:?} ({total_num_groups}) groups"
+ );
+ }
+ Ok(Self {
+ total_num_groups,
+ indices: Some(indices),
+ block_size,
+ })
+ }
+
+ /// Returns the group count against which this selection was constructed.
+ pub fn total_num_groups(self) -> usize {
+ self.total_num_groups
+ }
+
+ /// Ensures this selection is being applied to the same number of groups
+ /// against which it was constructed.
+ ///
+ /// Preserving-read implementations should call this method with their
+ /// stored group count before using [`Self::iter`]. This check is `O(1)`;
+ /// the selected indices were already checked by
[`Self::try_from_indices`].
+ pub fn validate_num_groups(self, actual_num_groups: usize) -> Result<()> {
+ if actual_num_groups != self.total_num_groups {
+ return exec_err!(
+ "Group selection was constructed for {} groups but applied to
{actual_num_groups} groups",
+ self.total_num_groups
+ );
+ }
+ Ok(())
+ }
+
+ /// Returns the number of selected groups.
+ pub fn len(self) -> usize {
+ self.indices
+ .map_or(self.total_num_groups, |indices| indices.len())
+ }
+
+ /// Returns `true` if no groups are selected.
+ pub fn is_empty(self) -> bool {
+ self.len() == 0
+ }
+
+ /// Returns the selected group indices in output order.
+ pub fn iter(self) -> impl Iterator<Item = BlocksIndex> + 'a {
+ let (all, indices): (_, &'a [BlocksIndex]) = match self.indices {
+ None => (0..self.total_num_groups, &[]),
+ Some(indices) => (0..0, indices),
+ };
+ let block_size = self.block_size;
+ all.map(move |index| {
+ BlocksIndex::from_index_in_fixed_block_size(index, block_size)
+ })
+ .chain(indices.iter().copied())
+ }
+
+ #[doc(hidden)]
+ pub fn indices(&self) -> Option<&[BlocksIndex]> {
+ self.indices
+ }
+}
+
+/// Two `u32`s (8 bytes) rather than two `usize`s: an index is stored per
group in every blocked
+/// hash table entry, so its size decides the size of the tables. Both halves
fit `u32` by a wide
+/// margin, `index_in_block < block_size` (the batch size) and the number of
blocks is
+/// `groups / block_size`. The API keeps taking and returning `usize`.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Hash)]
+pub struct BlocksIndex {
+ block_index: u32,
+ index_in_block: u32,
+}
+
+impl BlocksIndex {
+ pub const ZERO: Self = Self {
+ // flat_index: 0,
+ block_index: 0,
+ index_in_block: 0,
+ };
+ pub const MAX: Self = Self {
+ // flat_index: usize::MAX,
+ block_index: u32::MAX,
+ index_in_block: u32::MAX,
+ };
+
+ pub fn new(block_index: usize, index_in_block: usize) -> Self {
+ Self {
+ block_index: Self::narrow(block_index),
+ index_in_block: Self::narrow(index_in_block),
+ }
+ }
+
+ #[inline(always)]
+ fn narrow(value: usize) -> u32 {
+ debug_assert!(
+ u32::try_from(value).is_ok(),
+ "block index part {value} does not fit in u32"
+ );
+ value as u32
+ }
+
+ #[inline(always)]
+ pub fn new_in_first_block(index_in_block: usize) -> Self {
+ // Implementation note:
+ // not having From<usize> that will do this instead even when it will
be more convenient
+ // so we can later change the layout to be a single usize with bit
shifts
+ Self::new(0, index_in_block)
+ // Self {
+ // flat_index: index_in_block,
+ // }
+ }
+
+ #[inline(always)]
+ pub fn from_index_in_fixed_block_size(index: usize, block_size: usize) ->
Self {
+ Self::new(index / block_size, index % block_size)
+ }
+
+ #[inline(always)]
+ pub fn into_index_in_fixed_block_size(self, block_size: usize) -> usize {
+ self.block_index as usize * block_size + self.index_in_block as usize
+ }
+
+ #[inline(always)]
+ pub fn block_index(&self) -> usize {
+ self.block_index as usize
+ }
+
+ #[inline(always)]
+ pub fn index_in_block(&self) -> usize {
+ self.index_in_block as usize
+ }
+
+ pub fn sub_flat_checked(self, rhs_flat: usize, block_size: usize) ->
Option<Self> {
+ let self_flat = self.into_index_in_fixed_block_size(block_size);
+ self_flat
+ .checked_sub(rhs_flat)
+ .map(|v| Self::from_index_in_fixed_block_size(v, block_size))
+ }
+
+ pub fn sub_flat(self, rhs_flat: usize, block_size: usize) -> Self {
+ Self::from_index_in_fixed_block_size(
+ self.into_index_in_fixed_block_size(block_size) - rhs_flat,
+ block_size,
+ )
+ }
+
+ pub fn gte_flat(self, rhs_flat: usize, block_size: usize) -> bool {
+ self.into_index_in_fixed_block_size(block_size) >= rhs_flat
+ }
+
+ pub fn prev_block(mut self) -> Self {
+ self.block_index -= 1;
+
+ self
+ }
+
+ pub fn prev_block_checked(mut self) -> Option<Self> {
+ self.block_index.checked_sub(1).map(|v| {
+ self.block_index = v;
+ self
+ })
+ }
+
+ pub fn add(self, n: usize, block_size: usize) -> BlocksIndex {
+ if self.index_in_block + n as u32 >= block_size as u32 {
+ let new_flat = self.into_index_in_fixed_block_size(block_size) + n;
+ Self::from_index_in_fixed_block_size(new_flat, block_size)
+ } else {
+ Self::new(self.block_index as usize, self.index_in_block as usize
+ n)
+ }
+ }
+
+ pub fn increment(mut self, block_size: usize) -> Self {
+ if self.index_in_block + 1 >= block_size as u32 {
+ self.block_index += 1;
+ self.index_in_block = 0;
+ } else {
+ self.index_in_block += 1;
+ }
+
+ self
+ }
+}
+
+impl PartialOrd for BlocksIndex {
+ fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+ Some(self.cmp(other))
+ }
+}
+
+impl Ord for BlocksIndex {
+ fn cmp(&self, other: &Self) -> Ordering {
+ self.block_index
+ .cmp(&other.block_index)
+ .then(self.index_in_block.cmp(&other.index_in_block))
+ }
+}
+
+/// Describes how many rows should be emitted during grouping.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum BlockedEmitTo {
+ All,
+ /// Emit next group
+ NextBlock,
+ /// Emit only the first `n` groups and shift all existing group
+ /// indexes down by `n`.
+ ///
+ /// For example, if `n=10`, group_index `0, 1, ... 9` are emitted
+ /// and group indexes `10, 11, 12, ...` become `0, 1, 2, ...`.
+ ///
+ /// Requirements:
+ /// 1. `n` must be smaller than block_size
+ /// 2. `n` is not 0
+ First(usize),
+}
+
+impl BlockedEmitTo {
+ pub fn into_emit_to(
+ self,
+ len: usize,
+ block_size: usize,
+ ) -> Result<impl Iterator<Item = EmitTo> + Clone> {
+ let mut to_emit = match self {
+ Self::All => len,
+ Self::NextBlock => len.min(block_size),
+ Self::First(n) => {
+ assert_ne_or_internal_err!(n, 0);
+ assert_or_internal_err!(
+ n <= len,
+ "n ({n}) must be less than or equal current length ({})",
+ len
+ );
+ assert_or_internal_err!(
+ n < block_size,
+ "n ({n}) must be less than current block size ({})",
+ block_size
+ );
+ n
+ }
+ };
+
+ let mut blocks = vec![];
+
+ let emitting_all = len == to_emit;
+
+ while to_emit > block_size {
+ blocks.push(EmitTo::First(block_size));
+ to_emit -= block_size
+ }
+
+ if to_emit > 0 && emitting_all {
+ blocks.push(EmitTo::All);
+ } else if to_emit > 0 {
+ blocks.push(EmitTo::First(to_emit));
+ }
+
+ Ok(blocks.into_iter())
+ }
+}
+
+pub trait BlockedGroupsAccumulator: Send + std::any::Any {
+ fn batch_size(&self) -> usize;
Review Comment:
nit: this is named `batch_size()`, but everything else
(`BlockedVec::block_size`, `BlockedGroupValues::block_size`) says block size.
Suggest `fn block_size(&self) -> usize;`, and also renaming
`BlockedAccumulatorArgs::batch_size`.
##########
datafusion/expr-common/src/blocked_groups_accumulator.rs:
##########
@@ -0,0 +1,398 @@
+// 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.
+
+//! Vectorized [`GroupsAccumulator`]
+
+use crate::groups_accumulator::EmitTo;
+use arrow::array::{ArrayRef, BooleanArray};
+use datafusion_common::{
+ Result, assert_ne_or_internal_err, assert_or_internal_err, exec_err,
not_impl_err,
+};
+use std::cmp::Ordering;
+
+/// Selects groups for a non-destructive grouped aggregation read.
+///
+/// Unlike [`BlockedEmitTo`], this selection does not remove groups or change
their
+/// indices. Selections created by [`Self::try_from_indices`] preserve the
+/// requested order and support duplicate indices.
+///
+/// A selection is validated once when it is constructed and can then be reused
+/// for the group values and accumulators participating in the same snapshot.
+/// Construct a new selection if the number or indexing of those groups
changes.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct BlockedGroupSelection<'a> {
+ block_size: usize,
+ total_num_groups: usize,
+ indices: Option<&'a [BlocksIndex]>,
+}
+
+impl<'a> BlockedGroupSelection<'a> {
+ /// Selects all `total_num_groups` groups in group-index order.
+ pub fn all(total_num_groups: usize, block_size: usize) -> Self {
+ Self {
+ total_num_groups,
+ indices: None,
+ block_size,
+ }
+ }
+
+ /// Selects groups in the order specified by `indices`.
+ ///
+ /// Returns an error if an index is not less than `total_num_groups`. Empty
+ /// selections are valid, and duplicate indices are preserved.
+ pub fn try_from_indices(
+ indices: &'a [BlocksIndex],
+ total_num_groups: usize,
+ block_size: usize,
+ ) -> Result<Self> {
+ let total_num_groups_parsed =
+ BlocksIndex::from_index_in_fixed_block_size(total_num_groups,
block_size);
+ if let Some(index) = indices
+ .iter()
+ .find(|&&index| index >= total_num_groups_parsed)
+ {
+ return exec_err!(
+ "Group index {index:?} is out of bounds for
{total_num_groups_parsed:?} ({total_num_groups}) groups"
+ );
+ }
+ Ok(Self {
+ total_num_groups,
+ indices: Some(indices),
+ block_size,
+ })
+ }
+
+ /// Returns the group count against which this selection was constructed.
+ pub fn total_num_groups(self) -> usize {
+ self.total_num_groups
+ }
+
+ /// Ensures this selection is being applied to the same number of groups
+ /// against which it was constructed.
+ ///
+ /// Preserving-read implementations should call this method with their
+ /// stored group count before using [`Self::iter`]. This check is `O(1)`;
+ /// the selected indices were already checked by
[`Self::try_from_indices`].
+ pub fn validate_num_groups(self, actual_num_groups: usize) -> Result<()> {
+ if actual_num_groups != self.total_num_groups {
+ return exec_err!(
+ "Group selection was constructed for {} groups but applied to
{actual_num_groups} groups",
+ self.total_num_groups
+ );
+ }
+ Ok(())
+ }
+
+ /// Returns the number of selected groups.
+ pub fn len(self) -> usize {
+ self.indices
+ .map_or(self.total_num_groups, |indices| indices.len())
+ }
+
+ /// Returns `true` if no groups are selected.
+ pub fn is_empty(self) -> bool {
+ self.len() == 0
+ }
+
+ /// Returns the selected group indices in output order.
+ pub fn iter(self) -> impl Iterator<Item = BlocksIndex> + 'a {
+ let (all, indices): (_, &'a [BlocksIndex]) = match self.indices {
+ None => (0..self.total_num_groups, &[]),
+ Some(indices) => (0..0, indices),
+ };
+ let block_size = self.block_size;
+ all.map(move |index| {
+ BlocksIndex::from_index_in_fixed_block_size(index, block_size)
+ })
+ .chain(indices.iter().copied())
+ }
+
+ #[doc(hidden)]
+ pub fn indices(&self) -> Option<&[BlocksIndex]> {
+ self.indices
+ }
+}
+
+/// Two `u32`s (8 bytes) rather than two `usize`s: an index is stored per
group in every blocked
+/// hash table entry, so its size decides the size of the tables. Both halves
fit `u32` by a wide
+/// margin, `index_in_block < block_size` (the batch size) and the number of
blocks is
+/// `groups / block_size`. The API keeps taking and returning `usize`.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Hash)]
+pub struct BlocksIndex {
+ block_index: u32,
+ index_in_block: u32,
+}
+
+impl BlocksIndex {
+ pub const ZERO: Self = Self {
+ // flat_index: 0,
+ block_index: 0,
+ index_in_block: 0,
+ };
+ pub const MAX: Self = Self {
+ // flat_index: usize::MAX,
+ block_index: u32::MAX,
+ index_in_block: u32::MAX,
+ };
+
+ pub fn new(block_index: usize, index_in_block: usize) -> Self {
+ Self {
+ block_index: Self::narrow(block_index),
+ index_in_block: Self::narrow(index_in_block),
+ }
+ }
+
+ #[inline(always)]
+ fn narrow(value: usize) -> u32 {
+ debug_assert!(
+ u32::try_from(value).is_ok(),
+ "block index part {value} does not fit in u32"
+ );
+ value as u32
+ }
+
+ #[inline(always)]
+ pub fn new_in_first_block(index_in_block: usize) -> Self {
+ // Implementation note:
+ // not having From<usize> that will do this instead even when it will
be more convenient
+ // so we can later change the layout to be a single usize with bit
shifts
+ Self::new(0, index_in_block)
+ // Self {
+ // flat_index: index_in_block,
+ // }
+ }
+
+ #[inline(always)]
+ pub fn from_index_in_fixed_block_size(index: usize, block_size: usize) ->
Self {
+ Self::new(index / block_size, index % block_size)
+ }
+
+ #[inline(always)]
+ pub fn into_index_in_fixed_block_size(self, block_size: usize) -> usize {
+ self.block_index as usize * block_size + self.index_in_block as usize
+ }
+
+ #[inline(always)]
+ pub fn block_index(&self) -> usize {
+ self.block_index as usize
+ }
+
+ #[inline(always)]
+ pub fn index_in_block(&self) -> usize {
+ self.index_in_block as usize
+ }
+
+ pub fn sub_flat_checked(self, rhs_flat: usize, block_size: usize) ->
Option<Self> {
+ let self_flat = self.into_index_in_fixed_block_size(block_size);
+ self_flat
+ .checked_sub(rhs_flat)
+ .map(|v| Self::from_index_in_fixed_block_size(v, block_size))
+ }
+
+ pub fn sub_flat(self, rhs_flat: usize, block_size: usize) -> Self {
+ Self::from_index_in_fixed_block_size(
+ self.into_index_in_fixed_block_size(block_size) - rhs_flat,
+ block_size,
+ )
+ }
+
+ pub fn gte_flat(self, rhs_flat: usize, block_size: usize) -> bool {
+ self.into_index_in_fixed_block_size(block_size) >= rhs_flat
+ }
+
+ pub fn prev_block(mut self) -> Self {
+ self.block_index -= 1;
+
+ self
+ }
+
+ pub fn prev_block_checked(mut self) -> Option<Self> {
+ self.block_index.checked_sub(1).map(|v| {
+ self.block_index = v;
+ self
+ })
+ }
+
+ pub fn add(self, n: usize, block_size: usize) -> BlocksIndex {
+ if self.index_in_block + n as u32 >= block_size as u32 {
+ let new_flat = self.into_index_in_fixed_block_size(block_size) + n;
+ Self::from_index_in_fixed_block_size(new_flat, block_size)
+ } else {
+ Self::new(self.block_index as usize, self.index_in_block as usize
+ n)
+ }
+ }
+
+ pub fn increment(mut self, block_size: usize) -> Self {
+ if self.index_in_block + 1 >= block_size as u32 {
+ self.block_index += 1;
+ self.index_in_block = 0;
+ } else {
+ self.index_in_block += 1;
+ }
+
+ self
+ }
+}
+
+impl PartialOrd for BlocksIndex {
+ fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+ Some(self.cmp(other))
+ }
+}
+
+impl Ord for BlocksIndex {
+ fn cmp(&self, other: &Self) -> Ordering {
+ self.block_index
+ .cmp(&other.block_index)
+ .then(self.index_in_block.cmp(&other.index_in_block))
+ }
+}
+
+/// Describes how many rows should be emitted during grouping.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum BlockedEmitTo {
+ All,
+ /// Emit next group
+ NextBlock,
+ /// Emit only the first `n` groups and shift all existing group
+ /// indexes down by `n`.
+ ///
+ /// For example, if `n=10`, group_index `0, 1, ... 9` are emitted
+ /// and group indexes `10, 11, 12, ...` become `0, 1, 2, ...`.
+ ///
+ /// Requirements:
+ /// 1. `n` must be smaller than block_size
+ /// 2. `n` is not 0
+ First(usize),
+}
+
+impl BlockedEmitTo {
+ pub fn into_emit_to(
+ self,
+ len: usize,
+ block_size: usize,
+ ) -> Result<impl Iterator<Item = EmitTo> + Clone> {
+ let mut to_emit = match self {
+ Self::All => len,
+ Self::NextBlock => len.min(block_size),
+ Self::First(n) => {
+ assert_ne_or_internal_err!(n, 0);
+ assert_or_internal_err!(
+ n <= len,
+ "n ({n}) must be less than or equal current length ({})",
+ len
+ );
+ assert_or_internal_err!(
+ n < block_size,
+ "n ({n}) must be less than current block size ({})",
+ block_size
+ );
+ n
+ }
+ };
+
+ let mut blocks = vec![];
+
+ let emitting_all = len == to_emit;
+
+ while to_emit > block_size {
+ blocks.push(EmitTo::First(block_size));
+ to_emit -= block_size
+ }
+
+ if to_emit > 0 && emitting_all {
+ blocks.push(EmitTo::All);
+ } else if to_emit > 0 {
+ blocks.push(EmitTo::First(to_emit));
+ }
+
+ Ok(blocks.into_iter())
+ }
+}
+
+pub trait BlockedGroupsAccumulator: Send + std::any::Any {
+ fn batch_size(&self) -> usize;
+
+ fn update_batch(
+ &mut self,
+ values: &[ArrayRef],
+ // Get blocks index rather than flat index to avoid extra computation
when inserting and searching (BlocksIndex can be implemented as flat usize)
+ group_indices: &[BlocksIndex],
+ opt_filter: Option<&BooleanArray>,
+ total_num_groups: usize,
+ ) -> Result<()>;
+
+ /// Same as `GroupsAccumulator::evaluate` but returning vec of blocks and
different EmitTo
+ /// For
+ /// - `BlockedEmitTo::All` it should return `Vec<Block>`
+ /// - `BlockedEmitTo::NextBlock` it should return single item vector with
the block or empty vec in case of no blocks
+ /// - `BlockedEmitTo::First(n)` it should return single item vector with
the first n rows in the first block. n must be smaller than block size and
length
+ ///
+ fn evaluate(&mut self, emit_to: BlockedEmitTo) -> Result<Vec<ArrayRef>>;
Review Comment:
Return shapes:
- `NextBlock` yields 0 or 1 blocks, so returning a `Vec` hides that;
`Option<ArrayRef>` (or a dedicated `evaluate_next_block`) would say so.
- `evaluate_preserving` (L348) and `state_preserving` return one flat array
rather than blocks, which is inconsistent with `evaluate` and `state`.
- If the plan is an iterator for emit All (the TODO in
`group_values/blocked.rs`), it's cheaper to put it in the trait now than to
change it later.
##########
datafusion/expr-common/src/blocked_groups_accumulator.rs:
##########
@@ -0,0 +1,398 @@
+// 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.
+
+//! Vectorized [`GroupsAccumulator`]
+
+use crate::groups_accumulator::EmitTo;
+use arrow::array::{ArrayRef, BooleanArray};
+use datafusion_common::{
+ Result, assert_ne_or_internal_err, assert_or_internal_err, exec_err,
not_impl_err,
+};
+use std::cmp::Ordering;
+
+/// Selects groups for a non-destructive grouped aggregation read.
+///
+/// Unlike [`BlockedEmitTo`], this selection does not remove groups or change
their
+/// indices. Selections created by [`Self::try_from_indices`] preserve the
+/// requested order and support duplicate indices.
+///
+/// A selection is validated once when it is constructed and can then be reused
+/// for the group values and accumulators participating in the same snapshot.
+/// Construct a new selection if the number or indexing of those groups
changes.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct BlockedGroupSelection<'a> {
+ block_size: usize,
+ total_num_groups: usize,
+ indices: Option<&'a [BlocksIndex]>,
+}
+
+impl<'a> BlockedGroupSelection<'a> {
+ /// Selects all `total_num_groups` groups in group-index order.
+ pub fn all(total_num_groups: usize, block_size: usize) -> Self {
+ Self {
+ total_num_groups,
+ indices: None,
+ block_size,
+ }
+ }
+
+ /// Selects groups in the order specified by `indices`.
+ ///
+ /// Returns an error if an index is not less than `total_num_groups`. Empty
+ /// selections are valid, and duplicate indices are preserved.
+ pub fn try_from_indices(
+ indices: &'a [BlocksIndex],
+ total_num_groups: usize,
+ block_size: usize,
+ ) -> Result<Self> {
+ let total_num_groups_parsed =
+ BlocksIndex::from_index_in_fixed_block_size(total_num_groups,
block_size);
+ if let Some(index) = indices
+ .iter()
+ .find(|&&index| index >= total_num_groups_parsed)
+ {
+ return exec_err!(
+ "Group index {index:?} is out of bounds for
{total_num_groups_parsed:?} ({total_num_groups}) groups"
+ );
+ }
+ Ok(Self {
+ total_num_groups,
+ indices: Some(indices),
+ block_size,
+ })
+ }
+
+ /// Returns the group count against which this selection was constructed.
+ pub fn total_num_groups(self) -> usize {
+ self.total_num_groups
+ }
+
+ /// Ensures this selection is being applied to the same number of groups
+ /// against which it was constructed.
+ ///
+ /// Preserving-read implementations should call this method with their
+ /// stored group count before using [`Self::iter`]. This check is `O(1)`;
+ /// the selected indices were already checked by
[`Self::try_from_indices`].
+ pub fn validate_num_groups(self, actual_num_groups: usize) -> Result<()> {
+ if actual_num_groups != self.total_num_groups {
+ return exec_err!(
+ "Group selection was constructed for {} groups but applied to
{actual_num_groups} groups",
+ self.total_num_groups
+ );
+ }
+ Ok(())
+ }
+
+ /// Returns the number of selected groups.
+ pub fn len(self) -> usize {
+ self.indices
+ .map_or(self.total_num_groups, |indices| indices.len())
+ }
+
+ /// Returns `true` if no groups are selected.
+ pub fn is_empty(self) -> bool {
+ self.len() == 0
+ }
+
+ /// Returns the selected group indices in output order.
+ pub fn iter(self) -> impl Iterator<Item = BlocksIndex> + 'a {
+ let (all, indices): (_, &'a [BlocksIndex]) = match self.indices {
+ None => (0..self.total_num_groups, &[]),
+ Some(indices) => (0..0, indices),
+ };
+ let block_size = self.block_size;
+ all.map(move |index| {
+ BlocksIndex::from_index_in_fixed_block_size(index, block_size)
+ })
+ .chain(indices.iter().copied())
+ }
+
+ #[doc(hidden)]
+ pub fn indices(&self) -> Option<&[BlocksIndex]> {
+ self.indices
+ }
+}
+
+/// Two `u32`s (8 bytes) rather than two `usize`s: an index is stored per
group in every blocked
+/// hash table entry, so its size decides the size of the tables. Both halves
fit `u32` by a wide
+/// margin, `index_in_block < block_size` (the batch size) and the number of
blocks is
+/// `groups / block_size`. The API keeps taking and returning `usize`.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Hash)]
+pub struct BlocksIndex {
Review Comment:
`BlocksIndex` puts a lot of public surface in `expr-common`: `sub_flat`,
`sub_flat_checked` (L199), `gte_flat`, `prev_block`, `add`, `increment`. Only
group-values implementations need these. Suggest keeping the public type to
`new`, `block_index`, `index_in_block` and the flat conversions, and moving the
rest next to their only user, or marking them `#[doc(hidden)]`.
##########
datafusion/expr-common/src/blocked_groups_accumulator.rs:
##########
@@ -0,0 +1,398 @@
+// 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.
+
+//! Vectorized [`GroupsAccumulator`]
+
+use crate::groups_accumulator::EmitTo;
+use arrow::array::{ArrayRef, BooleanArray};
+use datafusion_common::{
+ Result, assert_ne_or_internal_err, assert_or_internal_err, exec_err,
not_impl_err,
+};
+use std::cmp::Ordering;
+
+/// Selects groups for a non-destructive grouped aggregation read.
+///
+/// Unlike [`BlockedEmitTo`], this selection does not remove groups or change
their
+/// indices. Selections created by [`Self::try_from_indices`] preserve the
+/// requested order and support duplicate indices.
+///
+/// A selection is validated once when it is constructed and can then be reused
+/// for the group values and accumulators participating in the same snapshot.
+/// Construct a new selection if the number or indexing of those groups
changes.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct BlockedGroupSelection<'a> {
+ block_size: usize,
+ total_num_groups: usize,
+ indices: Option<&'a [BlocksIndex]>,
+}
+
+impl<'a> BlockedGroupSelection<'a> {
+ /// Selects all `total_num_groups` groups in group-index order.
+ pub fn all(total_num_groups: usize, block_size: usize) -> Self {
+ Self {
+ total_num_groups,
+ indices: None,
+ block_size,
+ }
+ }
+
+ /// Selects groups in the order specified by `indices`.
+ ///
+ /// Returns an error if an index is not less than `total_num_groups`. Empty
+ /// selections are valid, and duplicate indices are preserved.
+ pub fn try_from_indices(
+ indices: &'a [BlocksIndex],
+ total_num_groups: usize,
+ block_size: usize,
+ ) -> Result<Self> {
+ let total_num_groups_parsed =
+ BlocksIndex::from_index_in_fixed_block_size(total_num_groups,
block_size);
+ if let Some(index) = indices
+ .iter()
+ .find(|&&index| index >= total_num_groups_parsed)
+ {
+ return exec_err!(
+ "Group index {index:?} is out of bounds for
{total_num_groups_parsed:?} ({total_num_groups}) groups"
+ );
+ }
+ Ok(Self {
+ total_num_groups,
+ indices: Some(indices),
+ block_size,
+ })
+ }
+
+ /// Returns the group count against which this selection was constructed.
+ pub fn total_num_groups(self) -> usize {
+ self.total_num_groups
+ }
+
+ /// Ensures this selection is being applied to the same number of groups
+ /// against which it was constructed.
+ ///
+ /// Preserving-read implementations should call this method with their
+ /// stored group count before using [`Self::iter`]. This check is `O(1)`;
+ /// the selected indices were already checked by
[`Self::try_from_indices`].
+ pub fn validate_num_groups(self, actual_num_groups: usize) -> Result<()> {
+ if actual_num_groups != self.total_num_groups {
+ return exec_err!(
+ "Group selection was constructed for {} groups but applied to
{actual_num_groups} groups",
+ self.total_num_groups
+ );
+ }
+ Ok(())
+ }
+
+ /// Returns the number of selected groups.
+ pub fn len(self) -> usize {
+ self.indices
+ .map_or(self.total_num_groups, |indices| indices.len())
+ }
+
+ /// Returns `true` if no groups are selected.
+ pub fn is_empty(self) -> bool {
+ self.len() == 0
+ }
+
+ /// Returns the selected group indices in output order.
+ pub fn iter(self) -> impl Iterator<Item = BlocksIndex> + 'a {
+ let (all, indices): (_, &'a [BlocksIndex]) = match self.indices {
+ None => (0..self.total_num_groups, &[]),
+ Some(indices) => (0..0, indices),
+ };
+ let block_size = self.block_size;
+ all.map(move |index| {
+ BlocksIndex::from_index_in_fixed_block_size(index, block_size)
+ })
+ .chain(indices.iter().copied())
+ }
+
+ #[doc(hidden)]
+ pub fn indices(&self) -> Option<&[BlocksIndex]> {
+ self.indices
+ }
+}
+
+/// Two `u32`s (8 bytes) rather than two `usize`s: an index is stored per
group in every blocked
+/// hash table entry, so its size decides the size of the tables. Both halves
fit `u32` by a wide
+/// margin, `index_in_block < block_size` (the batch size) and the number of
blocks is
+/// `groups / block_size`. The API keeps taking and returning `usize`.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Hash)]
+pub struct BlocksIndex {
+ block_index: u32,
+ index_in_block: u32,
+}
+
+impl BlocksIndex {
+ pub const ZERO: Self = Self {
+ // flat_index: 0,
+ block_index: 0,
+ index_in_block: 0,
+ };
+ pub const MAX: Self = Self {
+ // flat_index: usize::MAX,
+ block_index: u32::MAX,
+ index_in_block: u32::MAX,
+ };
+
+ pub fn new(block_index: usize, index_in_block: usize) -> Self {
+ Self {
+ block_index: Self::narrow(block_index),
+ index_in_block: Self::narrow(index_in_block),
+ }
+ }
+
+ #[inline(always)]
+ fn narrow(value: usize) -> u32 {
+ debug_assert!(
+ u32::try_from(value).is_ok(),
+ "block index part {value} does not fit in u32"
+ );
+ value as u32
+ }
+
+ #[inline(always)]
+ pub fn new_in_first_block(index_in_block: usize) -> Self {
+ // Implementation note:
+ // not having From<usize> that will do this instead even when it will
be more convenient
+ // so we can later change the layout to be a single usize with bit
shifts
+ Self::new(0, index_in_block)
+ // Self {
+ // flat_index: index_in_block,
+ // }
+ }
+
+ #[inline(always)]
+ pub fn from_index_in_fixed_block_size(index: usize, block_size: usize) ->
Self {
+ Self::new(index / block_size, index % block_size)
+ }
+
+ #[inline(always)]
+ pub fn into_index_in_fixed_block_size(self, block_size: usize) -> usize {
+ self.block_index as usize * block_size + self.index_in_block as usize
+ }
+
+ #[inline(always)]
+ pub fn block_index(&self) -> usize {
+ self.block_index as usize
+ }
+
+ #[inline(always)]
+ pub fn index_in_block(&self) -> usize {
+ self.index_in_block as usize
+ }
+
+ pub fn sub_flat_checked(self, rhs_flat: usize, block_size: usize) ->
Option<Self> {
+ let self_flat = self.into_index_in_fixed_block_size(block_size);
+ self_flat
+ .checked_sub(rhs_flat)
+ .map(|v| Self::from_index_in_fixed_block_size(v, block_size))
+ }
+
+ pub fn sub_flat(self, rhs_flat: usize, block_size: usize) -> Self {
+ Self::from_index_in_fixed_block_size(
+ self.into_index_in_fixed_block_size(block_size) - rhs_flat,
+ block_size,
+ )
+ }
+
+ pub fn gte_flat(self, rhs_flat: usize, block_size: usize) -> bool {
+ self.into_index_in_fixed_block_size(block_size) >= rhs_flat
+ }
+
+ pub fn prev_block(mut self) -> Self {
+ self.block_index -= 1;
+
+ self
+ }
+
+ pub fn prev_block_checked(mut self) -> Option<Self> {
+ self.block_index.checked_sub(1).map(|v| {
+ self.block_index = v;
+ self
+ })
+ }
+
+ pub fn add(self, n: usize, block_size: usize) -> BlocksIndex {
+ if self.index_in_block + n as u32 >= block_size as u32 {
+ let new_flat = self.into_index_in_fixed_block_size(block_size) + n;
+ Self::from_index_in_fixed_block_size(new_flat, block_size)
+ } else {
+ Self::new(self.block_index as usize, self.index_in_block as usize
+ n)
+ }
+ }
+
+ pub fn increment(mut self, block_size: usize) -> Self {
+ if self.index_in_block + 1 >= block_size as u32 {
+ self.block_index += 1;
+ self.index_in_block = 0;
+ } else {
+ self.index_in_block += 1;
+ }
+
+ self
+ }
+}
+
+impl PartialOrd for BlocksIndex {
+ fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+ Some(self.cmp(other))
+ }
+}
+
+impl Ord for BlocksIndex {
+ fn cmp(&self, other: &Self) -> Ordering {
+ self.block_index
+ .cmp(&other.block_index)
+ .then(self.index_in_block.cmp(&other.index_in_block))
+ }
+}
+
+/// Describes how many rows should be emitted during grouping.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum BlockedEmitTo {
+ All,
+ /// Emit next group
+ NextBlock,
+ /// Emit only the first `n` groups and shift all existing group
+ /// indexes down by `n`.
+ ///
+ /// For example, if `n=10`, group_index `0, 1, ... 9` are emitted
+ /// and group indexes `10, 11, 12, ...` become `0, 1, 2, ...`.
+ ///
+ /// Requirements:
+ /// 1. `n` must be smaller than block_size
+ /// 2. `n` is not 0
+ First(usize),
Review Comment:
The `First(n)` precondition `n < block_size` leaks to callers:
`common_ordered.rs:481` already rebuilds `EmitTo::First(n)` as `NextBlock`s
plus a `First` remainder. Either allow any `n`, or define that decomposition
once, next to the enum:
```rs
impl BlockedEmitTo {
/// Splits a flat `EmitTo` into block-sized emits.
pub fn from_emit_to(emit_to: EmitTo, block_size: usize) -> Vec<Self> {
match emit_to {
EmitTo::All => vec![Self::All],
EmitTo::First(n) => {
let mut emits = vec![Self::NextBlock; n / block_size];
if n % block_size != 0 {
emits.push(Self::First(n % block_size));
}
emits
}
}
}
}
```
--
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]
