ariel-miculas commented on code in PR #15591: URL: https://github.com/apache/datafusion/pull/15591#discussion_r3748931416
########## datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/block_store/vec_block_store.rs: ########## @@ -0,0 +1,298 @@ +// 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. + +//! [`BlockStore`] wrapper specialized for stores whose blocks are `Vec<T>`. + +use std::fmt::Debug; +use std::iter; +use std::marker::PhantomData; +use std::ops::{Index, IndexMut}; + +use datafusion_common::utils::split_vec_min_alloc; +use datafusion_common::{Result, internal_datafusion_err, internal_err}; +use datafusion_expr_common::groups_accumulator::EmitTo; + +use crate::aggregate::groups_accumulator::block_store::{Block, BlockStore}; + +/// Thin wrapper around a [`BlockStore<Vec<T>>`] that adds vector-specific +/// emit semantics on top of the generic block store API. +/// +/// Emitting `EmitTo::First` requires vector-specific split semantics, so it +/// lives on this wrapper rather than on [`BlockStore`] itself. The wrapper +/// re-exposes every [`BlockStore`] method by delegation, plus an [`emit`] +/// method implemented purely via [`BlockStore::push_block`] / +/// [`BlockStore::pop_block`] so it works uniformly over flat and blocked +/// storage. +/// +/// [`emit`]: VecBlockStore::emit +#[derive(Debug)] +pub struct VecBlockStore<T, S> +where + T: Clone + Debug, + S: BlockStore<Vec<T>>, +{ + inner: S, + _phantom: PhantomData<T>, +} + +impl<T, S> VecBlockStore<T, S> +where + T: Clone + Debug, + S: BlockStore<Vec<T>>, +{ + /// Wrap an existing block store. + pub fn new(inner: S) -> Self { + Self { + inner, + _phantom: PhantomData, + } + } + + // ---- BlockStore method delegation ---------------------------------- + pub fn reserve_blocks(&mut self) { + self.inner.reserve_blocks(); + } + + pub fn resize(&mut self, total_num_groups: usize, default_value: T) { + self.inner.resize(total_num_groups, default_value); + } + + pub fn num_blocks(&self) -> usize { + self.inner.num_blocks() + } + + pub fn block_size(&self) -> Option<usize> { + self.inner.block_size() + } + + pub fn clear(&mut self) { + self.inner.clear(); + } + + // ---- Emit ---------------------------------------------------------- + + /// Emit values according to `emit_to`, expressed only in terms of + /// [`BlockStore::push_block`] and [`BlockStore::pop_block`]. + /// + /// - [`EmitTo::All`]: drains every block via repeated `pop_block` and + /// concatenates the results into a single `Vec<T>`. + /// - [`EmitTo::First`]`(n)`: pops the first block, splits off the first + /// `n` values, and pushes the remainder back. Only meaningful when the + /// first block holds at least `n` values (true for flat storage); the + /// call returns an internal error otherwise. + /// - [`EmitTo::NextBlock`]: pops a single block. + pub fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<T>> { + match emit_to { + EmitTo::All => self.inner.pop_block().ok_or_else(|| { Review Comment: this doesn't seem to match the above description: > [`EmitTo::All`]: drains every block via repeated `pop_block` and concatenates the results into a single `Vec<T>`. -- 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]
