use std::borrow::Cow;
use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, OnceLock, TryLockError};
use rudb_common::{
Error, Field, LogicalType, Memory, PhysicalType, Reservation, Result, Session, Stage, Value,
stage,
};
use rudb_kernels::{
Accumulator, NOWHERE, finish_run, group_tally, is_true, settle_extremes, update_general,
update_runs, update_shared_runs, update_tallied, whole_answers,
};
use rudb_pipeline::{Lease, Progress, Sink};
use rudb_plan::{Expr, ExprRef, Plan, Slice};
use rudb_vector::{Chunk, Data, Form, Selection, VECTOR_SIZE, Validity, Vector};
use crate::blocks::{self, Blocks};
use crate::buffer::Buffered;
use crate::group_count;
use crate::group_distinct;
use crate::group_mixed;
use crate::group_ranged;
use crate::key::{BigIntSet, Key, RowSet, mix, spread};
use crate::pairs::{self, together};
use crate::places::Places;
use crate::prepared::{Prepared, Scratch};
use crate::rows;
use crate::schema::Schema;
use crate::signed::SignedBlock;
use crate::spill::{Reader, Spill};
use crate::table::{Origin, Probe, Table, Walk, held_at, slot_at};
#[derive(Debug, Clone)]
struct Call {
name: String,
args: Vec<ExprRef>,
distinct: bool,
filter: Option<ExprRef>,
returns: LogicalType,
affine: Option<(usize, i64)>,
reads_total: Option<usize>,
repeats: Option<usize>,
}
impl Call {
fn folds(&self) -> bool {
self.affine.is_none() && self.reads_total.is_none() && self.repeats.is_none()
}
fn state_of(&self, at: usize) -> usize {
self.repeats.unwrap_or(at)
}
fn finishes_plainly(&self) -> bool {
self.affine.is_none() && self.reads_total.is_none()
}
}
fn integer_constant(plan: &Plan, expr: ExprRef) -> Option<i64> {
let Expr::Constant(value) = *plan.expr(expr) else { return None };
let Value::Integer(offset) = *plan.value(value) else { return None };
Some(i64::from(offset))
}
fn mark_affine_sums(plan: &Plan, calls: &mut [Call]) {
for at in 0..calls.len() {
if calls[at].name != "sum" || calls[at].distinct || calls[at].filter.is_some() {
continue;
}
let [argument] = calls[at].args.as_slice() else { continue };
let Expr::Function { name, args } = *plan.expr(*argument) else { continue };
if plan.string(name) != "+" || plan.expr_type(*argument) != &LogicalType::Integer {
continue;
}
let [left, right] = plan.expr_list(args) else { continue };
let Expr::Cast { input: base, try_cast: false } = *plan.expr(*left) else { continue };
if plan.expr_type(base) != &LogicalType::SmallInt {
continue;
}
let Some(offset) = integer_constant(plan, *right) else { continue };
let source = (0..at).find(|&source| {
calls[source].name == "sum"
&& !calls[source].distinct
&& calls[source].filter.is_none()
&& calls[source].returns == LogicalType::HugeInt
&& calls[source].args.as_slice() == [base]
});
if let Some(source) = source {
calls[at].affine = Some((source, offset));
}
}
}
fn mark_sums_from_means(plan: &Plan, calls: &mut [Call]) {
let whole = |call: &Call| {
let [argument] = call.args.as_slice() else { return false };
let of = plan.expr_type(*argument);
of.is_integer() || matches!(of, LogicalType::Decimal { .. })
};
for at in 0..calls.len() {
if calls[at].name != "sum"
|| calls[at].distinct
|| calls[at].filter.is_some()
|| calls[at].affine.is_some()
|| !whole(&calls[at])
{
continue;
}
if calls.iter().any(|call| call.affine.is_some_and(|(source, _)| source == at)) {
continue;
}
calls[at].reads_total = (0..calls.len()).find(|&source| {
calls[source].name == "avg"
&& !calls[source].distinct
&& calls[source].filter.is_none()
&& calls[source].folds()
&& calls[source].args == calls[at].args
});
}
}
fn mark_repeated_calls(calls: &mut [Call]) {
for at in 0..calls.len() {
if calls[at].distinct || !calls[at].folds() {
continue;
}
let source_of_another = calls.iter().any(|call| {
call.affine.is_some_and(|(source, _)| source == at) || call.reads_total == Some(at)
});
if source_of_another {
continue;
}
calls[at].repeats = (0..at).find(|&source| {
let earlier = &calls[source];
earlier.folds()
&& !earlier.distinct
&& earlier.name == calls[at].name
&& earlier.args == calls[at].args
&& earlier.filter == calls[at].filter
&& earlier.returns == calls[at].returns
});
}
}
#[inline(never)]
fn sum_from_mean(held: &Accumulator, returns: &LogicalType) -> Result<Value> {
let Some((total, seen)) = held.exact_total() else {
return Err(Error::out_of_range("a total too large for an exact sum".to_string()));
};
Accumulator::sum_of(total, seen, returns)?.finish()
}
#[derive(Debug)]
pub(crate) struct Aggregate<'a> {
plan: &'a Plan,
keys: Vec<ExprRef>,
groups: Vec<ExprRef>,
constants: Vec<Option<Value>>,
calls: Vec<Call>,
template: Option<Vec<Accumulator>>,
inputs: Prepared,
schema: Schema,
alone: bool,
partition_from: usize,
sets: bool,
by_vector: Vec<bool>,
every: bool,
count_only: bool,
compact_numeric: bool,
distinct_count: bool,
mixed_numeric_distinct: bool,
radix_distinct_count: bool,
top_counts: Option<(usize, usize)>,
having_count: Option<(usize, i64)>,
max_groups: Option<usize>,
presize: Option<u64>,
reserve: bool,
span: Option<(i128, u64)>,
ends: Option<(i128, u64)>,
clustered: bool,
grouped: bool,
agreed: Mutex<Option<Agreed>>,
settled: AtomicBool,
memory: Memory,
built: Mutex<Built>,
merged: Vec<Mutex<Partition>>,
started: AtomicUsize,
locally: AtomicBool,
dense: OnceLock<DenseCount>,
fixed: OnceLock<FixedExchange>,
bigint_distinct: OnceLock<BigIntDistinctExchange>,
encoded_count: OnceLock<Option<EncodedExchange>>,
grouped_distinct: OnceLock<Option<group_distinct::Exchange>>,
mixed: OnceLock<group_mixed::Exchange>,
counted: OnceLock<group_count::Exchange>,
ranged: OnceLock<group_ranged::Exchange>,
out: Buffered,
}
#[derive(Debug)]
struct DenseCount {
dictionary: Arc<Vector>,
partitions: Vec<Mutex<DensePartition>>,
held: Mutex<Vec<Reservation>>,
}
#[derive(Debug, Default)]
struct DensePartition {
runs: Vec<Blocks<u32>>,
nulls: i64,
}
#[derive(Debug)]
struct FixedExchange {
keys: [LogicalType; 2],
partitions: Vec<Mutex<FixedRuns>>,
held: Mutex<Vec<Reservation>>,
}
#[derive(Debug)]
struct BigIntDistinctExchange {
partitions: Vec<Mutex<BigIntDistinctRuns>>,
held: Mutex<Vec<Reservation>>,
}
#[derive(Debug, Default)]
struct BigIntDistinctRuns {
runs: Vec<Vec<i64>>,
}
#[derive(Debug, Default)]
struct BigIntDistinctPartition {
rows: Vec<i64>,
}
impl BigIntDistinctPartition {
fn footprint(&self) -> usize {
self.rows.capacity() * size_of::<i64>()
}
}
#[derive(Debug)]
struct EncodedCountExchange<S = i64> {
dictionary: Arc<Vector>,
leading: Vec<LogicalType>,
dictionary_nulls: bool,
code_bits: u32,
partitions: Vec<Mutex<EncodedCountRuns<S>>>,
held: Mutex<Vec<Reservation>>,
}
impl<S: Second> EncodedCountExchange<S> {
fn new(
(dictionary, leading, dictionary_nulls, code_bits): (
Arc<Vector>,
Vec<LogicalType>,
bool,
u32,
),
) -> Self {
Self {
dictionary,
leading,
dictionary_nulls,
code_bits,
partitions: (0..RADIX_PARTITIONS)
.map(|_| Mutex::new(EncodedCountRuns::default()))
.collect(),
held: Mutex::new(Vec::new()),
}
}
fn take(&self, runs: &mut [EncodedCountRun<S>], memory: Reservation) -> Result<()> {
for (partition, rows) in runs.iter_mut().enumerate() {
if rows.is_empty() {
continue;
}
let run = std::mem::take(rows);
self.partitions[partition].lock().map_err(poisoned)?.runs.push(run);
}
self.held.lock().map_err(poisoned)?.push(memory);
Ok(())
}
}
#[derive(Debug, Default)]
struct EncodedCountRuns<S = i64> {
runs: Vec<EncodedCountRun<S>>,
}
#[derive(Debug, Default)]
struct FixedRuns {
runs: Vec<FixedRun>,
}
pub(crate) fn largest<T: Ord>(
groups: usize,
bound: usize,
count: impl Fn(usize) -> T,
) -> Vec<usize> {
let mut best: Vec<usize> = Vec::with_capacity(bound.min(groups));
if bound == 0 {
return best;
}
for slot in 0..groups {
let value = count(slot);
if best.len() == bound && best.last().is_some_and(|&last| count(last) >= value) {
continue;
}
let at = best.partition_point(|&kept| count(kept) >= value);
best.insert(at, slot);
best.truncate(bound);
}
best
}
const SLOT_BITS: u32 = 24;
const SLOT_MASK: u32 = (1 << SLOT_BITS) - 1;
const EMPTY_SLOT: u32 = SLOT_MASK;
const fn slot_tag(hash: u64) -> u32 {
(((hash >> SLOT_BITS) as u32) & 0xff) << SLOT_BITS
}
fn one_call(at: usize) -> u64 {
u32::try_from(at).ok().and_then(|at| 1_u64.checked_shl(at)).unwrap_or(0)
}
fn slot_runs_of(slots: &[usize], into: &mut Vec<(usize, usize)>, users: usize) -> bool {
into.clear();
let Some(&last) = slots.last() else {
return false;
};
let most = slots.len().saturating_mul(users) / RUN_ROWS;
into.reserve(most.min(slots.len()) + 1);
let mut row = 1;
while row < slots.len() {
let end = (row + RUN_BLOCK).min(slots.len());
let block = &slots[row..end];
let prior = &slots[row - 1..end - 1];
let mut starts = 0_u64;
for (at, (&slot, &before)) in block.iter().zip(prior).enumerate() {
starts |= u64::from(slot != before) << at;
}
if into.len() + starts.count_ones() as usize > most {
into.clear();
return false;
}
while starts != 0 {
let at = starts.trailing_zeros() as usize;
starts &= starts - 1;
into.push((prior[at], row + at));
}
row = end;
}
into.push((last, slots.len()));
true
}
const RUN_BLOCK: usize = 64;
const RUN_ROWS: usize = 8;
fn bucket_for(slot: usize, hash: u64, what: &'static str) -> Result<u32> {
let slot = u32::try_from(slot).ok().filter(|&slot| slot < SLOT_MASK);
let slot = slot.ok_or_else(|| Error::out_of_memory(what))?;
Ok(slot_tag(hash) | slot)
}
#[derive(Debug, Clone, Copy)]
struct EncodedCountRecord<S = i64> {
first: i64,
second: S,
hash: u32,
third: u32,
}
trait Second: Copy + PartialEq + std::fmt::Debug + Default + Send + 'static {
fn keys(self, third: u32, code_bits: u32) -> (i64, u32);
}
impl Second for i64 {
fn keys(self, third: u32, _: u32) -> (i64, u32) {
(self, third)
}
}
impl Second for () {
fn keys(self, third: u32, _: u32) -> (i64, u32) {
(0, third)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Default)]
struct Tucked;
impl Second for Tucked {
fn keys(self, third: u32, code_bits: u32) -> (i64, u32) {
let third = u64::from(third);
((third >> code_bits) as i64, (third & ((1 << code_bits) - 1)) as u32)
}
}
#[derive(Debug)]
enum EncodedExchange {
Two(EncodedCountExchange<()>),
Three { tucked: EncodedCountExchange<Tucked>, wide: EncodedCountExchange<i64> },
}
impl EncodedExchange {
fn first(&self) -> (&Arc<Vector>, &[LogicalType], bool) {
match self {
Self::Two(exchange) => {
(&exchange.dictionary, &exchange.leading, exchange.dictionary_nulls)
}
Self::Three { wide, .. } => (&wide.dictionary, &wide.leading, wide.dictionary_nulls),
}
}
}
#[derive(Debug)]
enum EncodedRecords {
Two(Vec<EncodedCountRun<()>>),
Three { tucked: Vec<EncodedCountRun<Tucked>>, wide: Vec<EncodedCountRun<i64>> },
}
impl EncodedRecords {
fn new(keys: usize) -> Self {
if keys == 2 { Self::Two(runs()) } else { Self::Three { tucked: runs(), wide: runs() } }
}
fn footprint(&self) -> usize {
match self {
Self::Two(runs) => runs.iter().map(EncodedCountRun::footprint).sum(),
Self::Three { tucked, wide } => {
tucked.iter().map(EncodedCountRun::footprint).sum::<usize>()
+ wide.iter().map(EncodedCountRun::footprint).sum::<usize>()
}
}
}
}
fn runs<S>() -> Vec<EncodedCountRun<S>> {
(0..RADIX_PARTITIONS).map(|_| EncodedCountRun::default()).collect()
}
trait Put {
fn put(&mut self, partition: usize, record: EncodedCountRecord, valid: u8, weight: u32);
}
impl Put for Vec<EncodedCountRun<()>> {
#[inline]
fn put(&mut self, partition: usize, record: EncodedCountRecord, valid: u8, weight: u32) {
let EncodedCountRecord { first, hash, third, .. } = record;
let record = EncodedCountRecord { first, second: (), hash, third };
self[partition].scatter(record, valid, weight);
}
}
struct Tucking<'a> {
tucked: &'a mut [EncodedCountRun<Tucked>],
wide: &'a mut [EncodedCountRun<i64>],
code_bits: u32,
limit: u64,
}
impl Put for Tucking<'_> {
#[inline]
fn put(&mut self, partition: usize, record: EncodedCountRecord, valid: u8, weight: u32) {
let EncodedCountRecord { first, second, hash, third } = record;
if valid == EncodedValid::ALL && (second as u64) < self.limit {
let third = third | ((second as u64) << self.code_bits) as u32;
self.tucked[partition].scatter(
EncodedCountRecord { first, second: Tucked, hash, third },
valid,
weight,
);
} else {
self.wide[partition].scatter(record, valid, weight);
}
}
}
fn code_bits(len: usize) -> u32 {
usize::BITS - len.saturating_sub(1).leading_zeros()
}
fn tuck_limit(code_bits: u32) -> u64 {
if code_bits >= u32::BITS { 0 } else { 1 << (u32::BITS - code_bits) }
}
#[derive(Debug, Default)]
struct EncodedCountPartition<S = i64> {
rows: Vec<EncodedCountRecord<S>>,
validity: Vec<u8>,
weights: Vec<u32>,
}
const COMPACT_FROM: usize = 4_096;
struct EncodedValid;
impl EncodedValid {
const FIRST: u8 = 1;
const SECOND: u8 = 2;
const THIRD: u8 = 4;
const ALL: u8 = Self::FIRST | Self::SECOND | Self::THIRD;
}
impl<S: Second> EncodedCountPartition<S> {
fn push(&mut self, row: EncodedCountRecord<S>, valid: u8) {
let keeping = !self.validity.is_empty() || valid != EncodedValid::ALL;
if keeping {
self.validity.resize(self.rows.len(), EncodedValid::ALL);
}
self.rows.push(row);
if keeping {
self.validity.push(valid);
}
if !self.weights.is_empty() {
self.weights.push(1);
}
}
fn push_weighted(&mut self, row: EncodedCountRecord<S>, valid: u8, weight: u32) {
let weighing = weight != 1 || !self.weights.is_empty();
self.push(row, valid);
if weighing {
self.weights.resize(self.rows.len(), 1);
*self.weights.last_mut().expect("a weight was just pushed") = weight;
}
}
#[inline]
fn weight(&self, at: usize) -> i64 {
self.weights.get(at).map_or(1, |&weight| i64::from(weight))
}
}
#[derive(Debug)]
struct EncodedCountRun<S = i64> {
groups: Blocks<EncodedCountRecord<S>>,
weights: Blocks<u32>,
group_validity: Vec<u8>,
pending: Blocks<EncodedCountRecord<S>>,
pending_validity: Vec<u8>,
pending_weights: Vec<u32>,
room: usize,
}
impl<S> Default for EncodedCountRun<S> {
fn default() -> Self {
Self {
groups: Blocks::default(),
weights: Blocks::default(),
group_validity: Vec::new(),
pending: Blocks::default(),
pending_validity: Vec::new(),
pending_weights: Vec::new(),
room: COMPACT_FROM,
}
}
}
impl<S: Second> EncodedCountRun<S> {
#[cfg(test)]
fn push(&mut self, row: EncodedCountRecord<S>, valid: u8) {
self.push_weighted(row, valid, 1);
}
fn push_weighted(&mut self, row: EncodedCountRecord<S>, valid: u8, weight: u32) {
let keeping = !self.pending_validity.is_empty() || valid != EncodedValid::ALL;
if keeping {
self.pending_validity.resize(self.pending.len(), EncodedValid::ALL);
}
let weighing = !self.pending_weights.is_empty() || weight != 1;
if weighing {
self.pending_weights.resize(self.pending.len(), 1);
}
self.pending.push(row);
if keeping {
self.pending_validity.push(valid);
}
if weighing {
self.pending_weights.push(weight);
}
}
#[inline]
fn scatter(&mut self, row: EncodedCountRecord<S>, valid: u8, weight: u32) {
if self.pending.len() >= self.room {
self.compact();
}
self.push_weighted(row, valid, weight);
}
#[inline]
fn pending_weight(&self, at: usize) -> u32 {
self.pending_weights.get(at).copied().unwrap_or(1)
}
fn is_empty(&self) -> bool {
self.groups.is_empty() && self.pending.is_empty()
}
fn len(&self) -> usize {
self.groups.len() + self.pending.len()
}
#[cold]
fn compact(&mut self) {
let len = self.len();
if blocks::locate(len).0 >= (SLOT_MASK >> blocks::WITHIN_BITS) as usize {
self.room = len;
return;
}
let capacity = len.saturating_mul(2).next_power_of_two();
let mask = capacity - 1;
let mut buckets = vec![EMPTY_SLOT; capacity];
let tag = |hash: u32| ((hash >> 18) & 0xff) << SLOT_BITS;
let place = |block: usize, within: usize| (block << blocks::WITHIN_BITS | within) as u32;
for (block, values) in self.groups.slices().enumerate() {
for (within, row) in values.iter().enumerate() {
let mut slot = row.hash as usize & mask;
while buckets[slot] != EMPTY_SLOT {
slot = (slot + 1) & mask;
}
buckets[slot] = tag(row.hash) | place(block, within);
}
}
let pending_valid = self.pending_validity.is_empty();
let all_valid = pending_valid && self.group_validity.is_empty();
if !all_valid {
self.group_validity.resize(self.groups.len(), EncodedValid::ALL);
}
let pending_validity = std::mem::take(&mut self.pending_validity);
let pending_weights = std::mem::take(&mut self.pending_weights);
let (mut block, mut within) = blocks::locate(self.groups.len());
let mut at = 0;
for values in std::mem::take(&mut self.pending).into_blocks() {
for &row in &values {
let valid = if pending_valid { EncodedValid::ALL } else { pending_validity[at] };
let weight = pending_weights.get(at).copied().unwrap_or(1);
at += 1;
let mut slot = row.hash as usize & mask;
let tagged = tag(row.hash);
loop {
let bucket = buckets[slot];
if bucket == EMPTY_SLOT {
buckets[slot] = tagged | place(block, within);
self.groups.push(row);
self.weights.push(weight);
if !all_valid {
self.group_validity.push(valid);
}
within += 1;
if within == blocks::size(block) {
block += 1;
within = 0;
}
break;
}
if bucket & !SLOT_MASK == tagged {
let held = ((bucket & SLOT_MASK) >> blocks::WITHIN_BITS) as usize;
let offset =
(bucket & SLOT_MASK) as usize & ((1 << blocks::WITHIN_BITS) - 1);
let other = self.groups.slot(held, offset);
if other.hash == row.hash
&& other.first == row.first
&& other.second == row.second
&& other.third == row.third
&& (all_valid
|| self.group_validity[blocks::start(held) + offset] == valid)
{
let sum = self.weights.slot_mut(held, offset);
if let Some(total) = sum.checked_add(weight) {
*sum = total;
break;
}
}
}
slot = (slot + 1) & mask;
}
}
}
let kept = self.groups.len();
let limit = if kept * 4 > len * 3 { len * 2 } else { len };
self.room = limit - kept;
}
fn footprint(&self) -> usize {
self.groups.footprint()
+ self.weights.footprint()
+ self.group_validity.capacity() * size_of::<u8>()
+ self.pending.footprint()
+ self.pending_validity.capacity() * size_of::<u8>()
+ self.pending_weights.capacity() * size_of::<u32>()
}
}
#[derive(Debug, Clone, Copy)]
struct FixedRecord {
first: i64,
second: i32,
sum: i16,
mean: i16,
}
#[derive(Debug, Default)]
struct FixedPartition {
rows: Vec<FixedRecord>,
validity: Vec<u8>,
}
impl FixedRecord {
const FIRST: u8 = 1;
const SECOND: u8 = 2;
const SUM: u8 = 4;
const MEAN: u8 = 8;
const ALL: u8 = Self::FIRST | Self::SECOND | Self::SUM | Self::MEAN;
}
fn fixed_hash(row: FixedRecord, valid: u8) -> u64 {
const NOTHING: u64 = 0x9e37_79b9_7f4a_7c15;
let first = if valid & FixedRecord::FIRST != 0 { row.first as u64 } else { NOTHING };
let second =
if valid & FixedRecord::SECOND != 0 { i64::from(row.second) as u64 } else { NOTHING };
spread(mix(mix(0, first), second))
}
impl FixedPartition {
fn push(&mut self, row: FixedRecord, valid: u8) {
let keeping = !self.validity.is_empty() || valid != FixedRecord::ALL;
if keeping {
self.validity.resize(self.rows.len(), FixedRecord::ALL);
}
self.rows.push(row);
if keeping {
self.validity.push(valid);
}
}
}
#[derive(Debug, Default)]
struct FixedRun {
rows: Blocks<FixedRecord>,
validity: Vec<u8>,
}
impl FixedRun {
fn push(&mut self, row: FixedRecord, valid: u8) {
let keeping = !self.validity.is_empty() || valid != FixedRecord::ALL;
if keeping {
self.validity.resize(self.rows.len(), FixedRecord::ALL);
}
self.rows.push(row);
if keeping {
self.validity.push(valid);
}
}
fn footprint(&self) -> usize {
self.rows.footprint() + self.validity.capacity() * size_of::<u8>()
}
}
#[derive(Debug)]
struct Built {
chunks: Vec<Chunk>,
held: Vec<Reservation>,
instances: usize,
partitioning: bool,
local: bool,
}
#[derive(Debug, Default)]
struct Partition {
table: Option<Building>,
carried: Option<Building>,
pending: Vec<Building>,
}
const RADIX_PARTITIONS: usize = 64;
const DENSE_PARTITIONS: usize = 4;
const SPARSE_DENSE: usize = 16;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum Share {
Whole,
Passing,
Partition,
Local,
}
impl Share {
fn of(self, groups: u64) -> Option<u64> {
match self {
Self::Whole => Some(groups),
Self::Passing => Some(groups.min(PARTITION_FROM as u64)),
Self::Partition => Some((groups / RADIX_PARTITIONS as u64).max(1)),
Self::Local => None,
}
}
fn before_the_split(self) -> bool {
matches!(self, Self::Whole | Self::Passing)
}
}
fn spare(memory: &Memory, room: u64) -> bool {
memory.limit().is_none_or(|limit| memory.used().saturating_add(room) <= limit / 2)
}
const PARTITION_FROM: usize = 4_096;
const FIXED_PARTITION_FROM: usize = 16_384;
const GATHER_ROWS: usize = 16_384;
const WINDOW_RATE: usize = 2;
const WINDOW_SLACK: usize = 4_096;
const LOCAL_CACHE: u64 = 16 * 1024 * 1024;
impl<'a> Aggregate<'a> {
#[must_use]
pub(crate) fn in_session(mut self, session: &Session) -> Self {
self.inputs = self.inputs.in_session(session);
self
}
pub(crate) fn new(
plan: &'a Plan,
input: &Schema,
index: u32,
groups: Slice,
aggregates: Slice,
memory: &Memory,
) -> Result<(Self, Buffered)> {
let input_schema = input.clone();
let groups: Vec<ExprRef> = plan.expr_list(groups).to_vec();
let constants: Vec<Option<Value>> = groups
.iter()
.map(|&group| match *plan.expr(group) {
Expr::Constant(value) => Some(plan.value(value).clone()),
_ => None,
})
.collect();
let keys: Vec<ExprRef> = groups
.iter()
.zip(&constants)
.filter_map(|(&group, value)| value.is_none().then_some(group))
.collect();
let mut calls = Vec::new();
for &reference in plan.expr_list(aggregates) {
let Expr::Aggregate { name, args, distinct, filter } = *plan.expr(reference) else {
return Err(Error::internal(format!(
"expression {reference} is in the aggregate list of an Aggregate and is not an aggregate"
)));
};
calls.push(Call {
name: plan.string(name).to_string(),
args: plan.expr_list(args).to_vec(),
distinct,
filter,
returns: plan.expr_type(reference).clone(),
affine: None,
reads_total: None,
repeats: None,
});
}
if groups.is_empty() {
mark_affine_sums(plan, &mut calls);
}
let mut fields = Vec::with_capacity(groups.len() + calls.len());
for (at, &group) in groups.iter().enumerate() {
fields.push(Field::new(
group_name(plan, group, &input_schema, at),
plan.expr_type(group).clone(),
));
}
for call in &calls {
fields.push(Field::new(call.name.clone(), call.returns.clone()));
}
let schema = Schema::numbered(fields, index);
let alone = groups.is_empty();
let compact_numeric = !alone
&& calls.len() == 3
&& calls[0].name == "count_star"
&& calls[0].args.is_empty()
&& calls[1].name == "sum"
&& calls[1].args.len() == 1
&& plan.expr_type(calls[1].args[0]) == &LogicalType::SmallInt
&& calls[1].returns == LogicalType::HugeInt
&& calls[2].name == "avg"
&& calls[2].args.len() == 1
&& plan.expr_type(calls[2].args[0]) == &LogicalType::SmallInt
&& calls[2].returns == LogicalType::Double
&& calls.iter().all(|call| !call.distinct && call.filter.is_none());
let distinct_count = !alone
&& calls.len() == 1
&& calls[0].name == "count"
&& calls[0].distinct
&& calls[0].args.len() == 1
&& plan.expr_type(calls[0].args[0]) == &LogicalType::BigInt
&& calls[0].filter.is_none();
let mixed_numeric_distinct = !alone
&& calls.len() == 4
&& calls[0].name == "sum"
&& !calls[0].distinct
&& calls[0].args.len() == 1
&& plan.expr_type(calls[0].args[0]) == &LogicalType::SmallInt
&& calls[0].returns == LogicalType::HugeInt
&& calls[1].name == "count_star"
&& !calls[1].distinct
&& calls[1].args.is_empty()
&& calls[1].returns == LogicalType::BigInt
&& calls[2].name == "avg"
&& !calls[2].distinct
&& calls[2].args.len() == 1
&& plan.expr_type(calls[2].args[0]) == &LogicalType::SmallInt
&& calls[2].returns == LogicalType::Double
&& calls[3].name == "count"
&& calls[3].distinct
&& calls[3].args.len() == 1
&& plan.expr_type(calls[3].args[0]) == &LogicalType::BigInt
&& calls[3].returns == LogicalType::BigInt
&& calls.iter().all(|call| call.filter.is_none());
let radix_distinct_count = alone
&& calls.len() == 1
&& calls[0].name == "count"
&& calls[0].distinct
&& calls[0].args.len() == 1
&& plan.expr_type(calls[0].args[0]) == &LogicalType::BigInt
&& calls[0].filter.is_none();
if !compact_numeric && !distinct_count && !mixed_numeric_distinct && !radix_distinct_count {
mark_sums_from_means(plan, &mut calls);
mark_repeated_calls(&mut calls);
}
let mut inputs = keys.clone();
for call in &calls {
if call.folds() {
inputs.extend_from_slice(&call.args);
}
inputs.extend(call.filter);
}
let inputs = Prepared::shared(plan, &inputs, &input_schema)?;
let by_vector: Vec<bool> =
calls.iter().map(|call| alone && !call.distinct && call.filter.is_none()).collect();
let template: Option<Vec<Accumulator>> =
calls.iter().map(|call| Accumulator::new(&call.name, &call.returns).ok()).collect();
let out = Buffered::new();
let partition_from = if keys.iter().all(|&key| fixed_width(plan.expr_type(key))) {
FIXED_PARTITION_FROM
} else {
PARTITION_FROM
};
let aggregate = Self {
plan,
keys,
constants,
alone,
partition_from,
sets: calls.iter().any(|call| call.distinct),
every: by_vector.iter().all(|&yes| yes),
count_only: !alone
&& calls.len() == 1
&& calls[0].name == "count_star"
&& !calls[0].distinct
&& calls[0].filter.is_none(),
compact_numeric,
distinct_count,
mixed_numeric_distinct,
radix_distinct_count,
top_counts: None,
having_count: None,
max_groups: None,
presize: None,
reserve: false,
span: None,
ends: None,
clustered: false,
grouped: false,
agreed: Mutex::new(None),
settled: AtomicBool::new(false),
by_vector,
groups,
calls,
template,
inputs,
schema,
memory: memory.clone(),
built: Mutex::new(Built {
chunks: Vec::new(),
held: Vec::new(),
instances: 0,
partitioning: false,
local: true,
}),
merged: (0..RADIX_PARTITIONS).map(|_| Mutex::new(Partition::default())).collect(),
started: AtomicUsize::new(0),
locally: AtomicBool::new(true),
dense: OnceLock::new(),
fixed: OnceLock::new(),
bigint_distinct: OnceLock::new(),
encoded_count: OnceLock::new(),
grouped_distinct: OnceLock::new(),
mixed: OnceLock::new(),
counted: OnceLock::new(),
ranged: OnceLock::new(),
out: out.clone(),
};
Ok((aggregate, out))
}
pub(crate) fn limit_groups(mut self, max_groups: usize) -> Self {
self.max_groups = Some(max_groups);
self
}
pub(crate) fn presize(mut self, groups: u64) -> Self {
self.presize = Some(groups);
self
}
pub(crate) fn reserved(mut self) -> Self {
self.reserve = true;
self
}
pub(crate) fn over_range(mut self, low: i128, values: u64) -> Self {
self.span = Some((low, values));
self
}
pub(crate) fn within(mut self, low: i128, values: u64) -> Self {
self.ends = Some((low, values));
self
}
pub(crate) fn clustered(mut self) -> Self {
self.clustered = true;
self
}
pub(crate) fn grouped(mut self) -> Self {
self.clustered = true;
self.grouped = true;
self
}
fn closes(&self) -> bool {
self.clustered
&& !self.alone
&& !self.sets
&& self.keys.len() == 1
&& self.max_groups.is_none()
&& !self.count_only
&& !self.radix_distinct_count
&& !self.fixed_top_count()
&& !self.encoded_top_count()
&& !self.grouped_distinct_top_count()
&& !self.mixed_top_count()
&& !self.counted_top_count()
}
fn closes_by_run(&self) -> bool {
self.closes()
&& !self.compact_numeric
&& self.having_count.is_none()
&& self.top_counts.is_none()
&& self.calls.iter().enumerate().all(|(at, call)| {
!call.distinct
&& call.filter.is_none()
&& call.folds()
&& !self.by_vector[at]
&& match (call.name.as_str(), call.args.as_slice()) {
("count_star", []) | ("count", [_]) => true,
("sum", [argument]) => {
whole_total(self.plan.expr_type(*argument), &call.returns)
}
_ => false,
}
})
}
fn open(
&self,
rows: &Rows,
single: &mut Option<Building>,
installed: &mut bool,
spreading: &mut Spreading,
own: &mut [Option<Building>],
folded: &mut u64,
) -> Result<()> {
*folded += rows.rows as u64;
let Some(table) = single else {
spreading.gathered += rows.rows;
spreading.pending.push(rows.settled()?.into_owned());
if spreading.gathered < GATHER_ROWS {
return Ok(());
}
return self.drain(*folded, spreading, own);
};
if let Some(error) = table.failure.take() {
return Err(error);
}
if let Some(limit) = self.max_groups
&& !self.alone
{
self.agree(rows, limit, table, installed)?;
}
let timing = stage::Timing::start(Stage::Fold);
let done = self.fold(rows, table, None, None);
timing.stop(0);
done?;
if !self.ought_to_partition(table) {
return Ok(());
}
let handing = single.take().expect("the table was there a moment ago");
self.begin_partitioning(spreading, own)?;
self.hand(handing, spreading, own)
}
fn shut(&self) -> Building {
let mut local = self.starting(Share::Local);
let types: Vec<_> = self.keys.iter().map(|&key| self.plan.expr_type(key).clone()).collect();
local.table = Table::new(&types);
local
}
#[must_use]
pub(crate) fn top_counts(mut self, bound: usize, call: usize) -> Self {
if self.count_only
|| self.compact_numeric
|| self.distinct_count
|| self.mixed_numeric_distinct
{
self.top_counts = Some((bound, 0));
} else if self.calls.get(call).is_some_and(|held| {
held.name == "count_star"
&& held.args.is_empty()
&& !held.distinct
&& held.filter.is_none()
}) {
self.top_counts = Some((bound, self.calls[call].state_of(call)));
}
self
}
#[must_use]
pub(crate) fn having_count(mut self, call: usize, minimum: i64) -> Self {
if self.calls.get(call).is_some_and(|held| {
held.name == "count_star"
&& held.args.is_empty()
&& !held.distinct
&& held.filter.is_none()
}) {
self.having_count = Some((self.calls[call].state_of(call), minimum));
}
self
}
pub(crate) fn schema(&self) -> &Schema {
&self.schema
}
fn read(&self, chunk: &Chunk, scratch: &mut Scratch) -> Result<Rows> {
let mut evaluated = Vec::new();
let Some(kept) = chunk.kept() else {
self.inputs.evaluate(chunk, scratch, &mut evaluated)?;
return self.rows_of(evaluated, chunk.len(), None);
};
if self.reads_marked() && self.inputs.evaluate(chunk, scratch, &mut evaluated).is_ok() {
return self.rows_of(evaluated, chunk.len(), Some(kept));
}
self.read(&chunk.clone().settled()?, scratch)
}
fn reads_marked(&self) -> bool {
!self.alone
&& !self.sets
&& !self.compact_numeric
&& !self.distinct_count
&& !self.mixed_numeric_distinct
&& !self.radix_distinct_count
&& !(self.count_only && self.keys.len() == 1)
&& self.top_counts.is_none()
&& self.max_groups.is_none()
&& !self.closes()
&& !self.by_vector.iter().any(|&yes| yes)
&& self.calls.iter().all(|call| !call.distinct && call.affine.is_none())
}
fn rows_of(
&self,
evaluated: Vec<Vector>,
rows: usize,
kept: Option<&Selection>,
) -> Result<Rows> {
let mut values = evaluated.into_iter();
let mut keys: Vec<Vector> = values.by_ref().take(self.keys.len()).collect();
let (rows, marked) = match kept {
Some(kept) => {
keys = Chunk::with_rows(keys, rows)?.select(kept)?.into_columns();
(kept.len(), Some((kept.clone(), rows)))
}
None => (rows, None),
};
if keys.len() > 1 {
for key in &mut keys {
if key.logical_type().is_integer()
&& key.data().is_none()
&& key.constant_value().is_none()
{
*key = key.opened()?;
}
}
}
let mut arguments = Vec::with_capacity(self.calls.len());
let mut filters = Vec::with_capacity(self.calls.len());
for call in &self.calls {
arguments.push(if call.folds() {
values.by_ref().take(call.args.len()).collect()
} else {
Vec::new()
});
filters
.push(call.filter.map(|_| values.next().expect("a prepared filter has a value")));
}
Ok(Rows { keys, arguments, filters, rows, marked })
}
fn fixed_top_count(&self) -> bool {
self.compact_numeric
&& self.top_counts.is_some()
&& self.constants.iter().all(Option::is_none)
&& self.keys.len() == 2
&& signed_key(self.plan.expr_type(self.keys[0]))
&& narrow_key(self.plan.expr_type(self.keys[1]))
}
fn encoded_top_count(&self) -> bool {
if !self.count_only
|| self.top_counts.is_none()
|| !self.constants.iter().all(Option::is_none)
{
return false;
}
let Some((last, leading)) = self.keys.split_last() else { return false };
(1..=2).contains(&leading.len())
&& self.plan.expr_type(*last) == &LogicalType::Varchar
&& leading.iter().all(|&key| signed_key(self.plan.expr_type(key)))
}
fn grouped_distinct_top_count(&self) -> bool {
self.distinct_count
&& self.top_counts.is_some()
&& self.constants.iter().all(Option::is_none)
&& (1..=2).contains(&self.keys.len())
&& self.keys.iter().all(|&key| {
matches!(
self.plan.expr_type(key),
LogicalType::TinyInt
| LogicalType::SmallInt
| LogicalType::Integer
| LogicalType::Varchar
)
})
}
fn counted_top_count(&self) -> bool {
self.count_only
&& self.top_counts.is_some()
&& self.span.is_none()
&& self.having_count.is_none()
&& self.max_groups.is_none()
&& !self.sets
&& self.constants.iter().all(Option::is_none)
&& self.keys.len() == 1
&& signed_key(self.plan.expr_type(self.keys[0]))
}
fn ranged_counts(&self) -> bool {
!self.alone
&& self.span.is_some()
&& self.top_counts.is_none()
&& self.having_count.is_none()
&& self.max_groups.is_none()
&& !self.sets
&& self.keys.len() == 1
&& self.constants.iter().all(Option::is_none)
&& signed_key(self.plan.expr_type(self.keys[0]))
&& self.calls.iter().all(|call| {
call.folds()
&& !call.distinct
&& call.filter.is_none()
&& call.returns == LogicalType::BigInt
&& match call.name.as_str() {
"count_star" => call.args.is_empty(),
"count" => call.args.len() == 1,
_ => false,
}
})
}
fn count_ranged(&self, rows: &Rows, local: &mut group_ranged::Local) -> Result<Progress> {
let rows = rows.settled()?;
let [key] = rows.keys.as_slice() else {
return Err(Error::internal("a ranged count received the wrong key width"));
};
let exchange = self.ranged.get_or_init(|| {
let (low, values) = self.span.expect("a ranged count has a range");
let calls = self
.calls
.iter()
.map(|call| {
if call.args.is_empty() {
group_ranged::Counted::Rows
} else {
group_ranged::Counted::Valid
}
})
.collect();
group_ranged::Exchange::new(
self.plan.expr_type(self.keys[0]).clone(),
i64::try_from(low).unwrap_or(i64::MIN),
usize::try_from(values).unwrap_or(0),
calls,
)
});
let arguments: Vec<Option<&Vector>> =
rows.arguments.iter().map(|call| call.first()).collect();
exchange.count(key, &arguments, rows.rows, local)?;
Ok(Progress::More)
}
fn mixed_top_count(&self) -> bool {
self.mixed_numeric_distinct
&& self.top_counts.is_some()
&& self.constants.iter().all(Option::is_none)
&& self.keys.len() == 1
&& self.plan.expr_type(self.keys[0]) == &LogicalType::Integer
}
fn buffer_fixed(
&self,
rows: &Rows,
partitions: &mut [FixedRun],
memory: &mut Reservation,
blocks: &mut FixedBlocks,
) -> Result<()> {
self.fixed.get_or_init(|| FixedExchange {
keys: [
self.plan.expr_type(self.keys[0]).clone(),
self.plan.expr_type(self.keys[1]).clone(),
],
partitions: (0..RADIX_PARTITIONS).map(|_| Mutex::new(FixedRuns::default())).collect(),
held: Mutex::new(Vec::new()),
});
let [first, second] = rows.keys.as_slice() else {
return Err(Error::internal("a fixed radix exchange received the wrong key width"));
};
let sum = rows.arguments[1].first().expect("SUM has one argument");
let mean = rows.arguments[2].first().expect("AVG has one argument");
let before = partitions.iter().map(FixedRun::footprint).sum::<usize>();
let shift = u64::BITS - RADIX_PARTITIONS.ilog2();
blocks.read(rows.rows, [first, second, sum, mean])?;
let [held_first, held_second, held_sum, held_mean] = blocks.cut(rows.rows)?;
let [null_first, null_second, null_sum, null_mean] = blocks.nulled();
for row in 0..rows.rows {
let mut valid = 0;
let first_value = if null_first && first.is_null_at(row) {
0
} else {
valid |= FixedRecord::FIRST;
held_first[row]
};
let second_value = if null_second && second.is_null_at(row) {
0
} else {
valid |= FixedRecord::SECOND;
i32::try_from(held_second[row])
.map_err(|_| Error::internal("a fixed second key is out of range"))?
};
let sum_value = if null_sum && sum.is_null_at(row) {
0
} else {
valid |= FixedRecord::SUM;
i16::try_from(held_sum[row])
.map_err(|_| Error::internal("a fixed SMALLINT sum is out of range"))?
};
let mean_value = if null_mean && mean.is_null_at(row) {
0
} else {
valid |= FixedRecord::MEAN;
i16::try_from(held_mean[row])
.map_err(|_| Error::internal("a fixed SMALLINT mean is out of range"))?
};
let record = FixedRecord {
first: first_value,
second: second_value,
sum: sum_value,
mean: mean_value,
};
let hash = fixed_hash(record, valid);
partitions[(hash >> shift) as usize].push(record, valid);
}
let after = partitions.iter().map(FixedRun::footprint).sum::<usize>();
memory.grow(width_of(after.saturating_sub(before)))
}
fn buffer_encoded_count(
&self,
rows: &Rows,
records: &mut EncodedRecords,
memory: &mut Reservation,
) -> Result<bool> {
let [.., third] = rows.keys.as_slice() else {
return Err(Error::internal("an encoded count exchange received no keys"));
};
let dictionary = third.stable_dictionary_parts();
let state = self.encoded_count.get_or_init(|| {
dictionary.as_ref().map(|(_, dictionary)| {
let leading: Vec<LogicalType> = self.keys[..self.keys.len() - 1]
.iter()
.map(|&key| self.plan.expr_type(key).clone())
.collect();
let nulls = dictionary.validity().has_nulls(dictionary.len())
|| !nulls_are_in_the_mask(dictionary);
let fresh = |code_bits| (Arc::clone(dictionary), leading.clone(), nulls, code_bits);
if self.keys.len() == 2 {
EncodedExchange::Two(EncodedCountExchange::new(fresh(u32::BITS)))
} else {
EncodedExchange::Three {
tucked: EncodedCountExchange::new(fresh(code_bits(dictionary.len()))),
wide: EncodedCountExchange::new(fresh(u32::BITS)),
}
}
})
});
let Some(state) = state else { return Ok(false) };
let before = records.footprint();
let buffered = match (state, &mut *records) {
(EncodedExchange::Two(_), EncodedRecords::Two(runs)) => {
self.scatter_encoded_count(rows, state, runs)
}
(
EncodedExchange::Three { tucked, .. },
EncodedRecords::Three { tucked: into, wide },
) => {
let mut tucking = Tucking {
tucked: into,
wide,
code_bits: tucked.code_bits,
limit: tuck_limit(tucked.code_bits),
};
self.scatter_encoded_count(rows, state, &mut tucking)
}
_ => Err(Error::internal("an encoded count exchange changed key width")),
};
buffered?;
let after = records.footprint();
memory.grow(width_of(after.saturating_sub(before)))?;
Ok(true)
}
fn scatter_encoded_count(
&self,
rows: &Rows,
state: &EncodedExchange,
partitions: &mut impl Put,
) -> Result<()> {
let (first, second, third) = match rows.keys.as_slice() {
[first, third] => (first, None, third),
[first, second, third] => (first, Some(second), third),
_ => {
return Err(Error::internal(
"an encoded count exchange received the wrong key width",
));
}
};
let (dictionary_of_state, leading, dictionary_nulls) = state.first();
let Some((codes, dictionary)) = third.stable_dictionary_parts() else {
return Err(Error::internal(
"an encoded count exchange changed from dictionary to flat strings",
));
};
if !Arc::ptr_eq(dictionary_of_state, dictionary) {
return Err(Error::internal(
"an encoded count exchange received two string code spaces",
));
}
if leading.len() + 1 != rows.keys.len() {
return Err(Error::internal("an encoded count exchange changed key width"));
}
let shift = u32::BITS - RADIX_PARTITIONS.ilog2();
const NOTHING: u64 = 0x9e37_79b9_7f4a_7c15;
let plain = Signed::of(first, rows.rows).zip(match second {
Some(second) => Signed::of(second, rows.rows).map(Some),
None => Some(None),
});
let plain = plain.filter(|_| {
!dictionary_nulls && third.len() >= rows.rows && !third.validity().has_nulls(rows.rows)
});
if let Some((first, second)) = plain {
let mut run: Option<(EncodedCountRecord, u32)> = None;
for (row, &third_code) in codes.iter().enumerate().take(rows.rows) {
if third_code as usize >= dictionary.len() {
return Err(Error::internal("an encoded string code is out of range"));
}
let first_value = first.at(row);
let second_value = second.map_or(0, |second| second.at(row));
if let Some((held, weight)) = &mut run {
if held.first == first_value
&& held.second == second_value
&& held.third == third_code
&& *weight < u32::MAX
{
*weight += 1;
continue;
}
partitions.put(
(held.hash >> shift) as usize,
*held,
EncodedValid::ALL,
*weight,
);
}
let wide = spread(mix(
mix(mix(0, first_value as u64), second_value as u64),
u64::from(third_code),
));
let hash = (wide ^ (wide >> 32)) as u32;
let record = EncodedCountRecord {
first: first_value,
second: second_value,
hash,
third: third_code,
};
run = Some((record, 1));
}
if let Some((held, weight)) = run {
partitions.put((held.hash >> shift) as usize, held, EncodedValid::ALL, weight);
}
return Ok(());
}
for (row, &third_code) in codes.iter().enumerate().take(rows.rows) {
let mut valid = if second.is_none() { EncodedValid::SECOND } else { 0 };
let first_value = if first.is_null_at(row) {
0
} else {
valid |= EncodedValid::FIRST;
i64::try_from(first.signed_at(row).ok_or_else(|| {
Error::internal("an encoded BIGINT key has no signed representation")
})?)
.map_err(|_| Error::internal("an encoded BIGINT key is out of range"))?
};
let second_value = if let Some(second) = second {
if second.is_null_at(row) {
0
} else {
valid |= EncodedValid::SECOND;
i64::try_from(second.signed_at(row).ok_or_else(|| {
Error::internal("an encoded BIGINT key has no signed representation")
})?)
.map_err(|_| Error::internal("an encoded BIGINT key is out of range"))?
}
} else {
0
};
let third_value = if third.is_null_at(row) {
0
} else {
valid |= EncodedValid::THIRD;
if third_code as usize >= dictionary.len() {
return Err(Error::internal("an encoded string code is out of range"));
}
third_code
};
let first_word =
if valid & EncodedValid::FIRST != 0 { first_value as u64 } else { NOTHING };
let second_word =
if valid & EncodedValid::SECOND != 0 { second_value as u64 } else { NOTHING };
let third_word =
if valid & EncodedValid::THIRD != 0 { u64::from(third_value) } else { NOTHING };
let wide = spread(mix(mix(mix(0, first_word), second_word), third_word));
let hash = (wide ^ (wide >> 32)) as u32;
partitions.put(
(hash >> shift) as usize,
EncodedCountRecord {
first: first_value,
second: second_value,
hash,
third: third_value,
},
valid,
1,
);
}
Ok(())
}
fn buffer_bigint_distinct(
&self,
rows: &Rows,
partitions: &mut [BigIntDistinctPartition],
memory: &mut Reservation,
) -> Result<()> {
self.bigint_distinct.get_or_init(|| BigIntDistinctExchange {
partitions: (0..RADIX_PARTITIONS)
.map(|_| Mutex::new(BigIntDistinctRuns::default()))
.collect(),
held: Mutex::new(Vec::new()),
});
let Some(column) = rows.arguments.first().and_then(|arguments| arguments.first()) else {
return Err(Error::internal("a BIGINT distinct exchange received no argument"));
};
let before = partitions.iter().map(BigIntDistinctPartition::footprint).sum::<usize>();
let shift = u64::BITS - RADIX_PARTITIONS.ilog2();
let flat = match column.data() {
Some(Data::Int64(values)) if !column.validity().has_nulls(rows.rows) => {
values.get(..rows.rows)
}
_ => None,
};
if let Some(values) = flat {
for &value in values {
scatter_bigint(partitions, shift, value);
}
} else {
for row in 0..rows.rows {
if column.is_null_at(row) {
continue;
}
let value = i64::try_from(column.signed_at(row).ok_or_else(|| {
Error::internal("a distinct BIGINT value has no signed representation")
})?)
.map_err(|_| Error::internal("a distinct BIGINT value is out of range"))?;
scatter_bigint(partitions, shift, value);
}
}
let after = partitions.iter().map(BigIntDistinctPartition::footprint).sum::<usize>();
memory.grow(width_of(after.saturating_sub(before)))
}
fn again(
&self,
file: &mut Spill,
carried: Option<Building>,
chunks: &mut Vec<Chunk>,
held: &mut Reservation,
) -> Result<Option<Spill>> {
let mut spilled = Spilled::new(file.read()?, self.spilled_types());
let mut local = self.partition();
if let Some(error) = local.failure.take() {
return Err(error);
}
if let Some(carried) = carried {
self.merge(carried, &mut local)?;
}
while let Some(rows) = spilled.next(self)? {
let timing = stage::Timing::start(Stage::Fold);
let folded = self.fold(&rows, &mut local, None, None);
timing.stop(0);
folded?;
}
self.finish(local, chunks, held)
}
fn start(&self) -> Building {
let held = self.alone || self.max_groups.is_some();
self.starting(if held { Share::Whole } else { Share::Passing })
}
fn partition(&self) -> Building {
self.starting(Share::Partition)
}
fn kept(&self) -> Building {
self.starting(Share::Local)
}
fn starting(&self, share: Share) -> Building {
let calls = self.calls.len();
let types: Vec<_> = self.keys.iter().map(|&key| self.plan.expr_type(key).clone()).collect();
let mut containers = self.memory.reservation();
let mut charged = 0;
let groups = self.presize.and_then(|groups| share.of(groups));
let groups = match groups {
Some(groups) if self.reserve => {
let room = Table::room(groups);
(spare(&self.memory, room) && containers.grow(room).is_ok()).then(|| {
charged = room;
groups
})
}
groups => groups,
};
let mut local = Building {
scratch: self.memory.reservation(),
containers,
charged,
charged_keys: 0,
table: {
let table = match groups {
Some(groups) => Table::with_groups(&types, groups),
None => Table::new(&types),
};
match (self.span.filter(|_| share.before_the_split()), types.as_slice()) {
(Some((low, values)), [ty]) => table.over_range(low, values, ty),
_ => table,
}
},
states: Vec::new(),
counts: Vec::new(),
compact: Vec::new(),
overflow: HashMap::new(),
seen: Vec::new(),
groups: 0,
given: vec![Key(Vec::new()); calls],
hashes: Vec::new(),
slots: Vec::new(),
walk: Walk::default(),
coded_on: Vec::new(),
coded_places: Vec::new(),
coded_values: crate::table::Widened::default(),
slot_runs: Vec::new(),
coded_spent: 0,
coded_read: 0,
coded_map: Places::default(),
same: Vec::new(),
leaders: Vec::new(),
leader_slots: Vec::new(),
kept: Vec::new(),
affine_rows: vec![0; calls],
over: None,
away: Vec::new(),
failure: None,
};
let window = self
.ends
.filter(|_| share.before_the_split() && !self.alone)
.and_then(|(low, values)| crate::table::seeded_window(&types, low, values));
if let Some((low, places)) = window {
local.coded_on.push(Origin::Window(low, places));
local.coded_map = Places::seeded(places);
}
if self.alone {
local.groups = 1;
if let Err(error) = self.fresh(&mut local.states, &mut local.counts, &mut local.compact)
{
local.failure = Some(error);
}
if self.sets {
self.fresh_seen(&mut local.seen);
}
}
local
}
fn fold(
&self,
seen_rows: &Rows,
local: &mut Building,
prehashed: Option<&[u64]>,
closed: Option<(usize, usize)>,
) -> Result<()> {
let Building {
scratch,
containers,
charged,
charged_keys,
table,
states,
counts,
compact,
overflow,
seen,
groups,
given,
hashes,
slots,
walk,
coded_on,
coded_places,
coded_values,
slot_runs,
coded_spent,
coded_read,
coded_map,
same,
leaders,
leader_slots,
kept,
affine_rows,
over,
away,
failure: _,
} = local;
let calls = self.calls.len();
let alone = self.alone;
let settled;
let seen_rows = if seen_rows.marked.is_some()
&& (over.is_some() || prehashed.is_some() || closed.is_some() || alone)
{
settled = seen_rows.settled()?;
&*settled
} else {
seen_rows
};
let Rows { keys, arguments, filters, rows: length, marked } = seen_rows;
let mut aside = 0;
for at in 0..calls {
if !self.calls[at].folds() {
continue;
}
if self.by_vector[at] {
states[at].update_run(&arguments[at], *length)?;
if self.calls.iter().any(|call| call.affine.is_some_and(|(source, _)| source == at))
{
affine_rows[at] +=
i64::try_from(arguments[at][0].validity().count_valid(*length))
.map_err(|_| Error::out_of_range("too many rows in an aggregate"))?;
}
}
}
if alone && self.every {
return Ok(());
}
slots.clear();
let direct = if alone
|| over.is_some()
|| closed.is_some()
|| !crate::table::fits_the_map(table.len(), *length)
{
coded_on.clear();
None
} else {
crate::table::coded_within(keys, *length, coded_on, Some(coded_values))
};
*coded_read = coded_read.saturating_add(*length);
let grown = direct.as_ref().and_then(|codes| codes.grows(coded_on, coded_map.len()));
let direct = direct.filter(|codes| {
codes.same_as(coded_on)
|| !codes.reads_values()
|| coded_spent.saturating_add(codes.combos() - grown.unwrap_or(0))
<= coded_read.saturating_mul(WINDOW_RATE).saturating_add(WINDOW_SLACK)
});
let mut runs_found = false;
if let Some(codes) = &direct {
if !codes.same_as(coded_on) {
if codes.reads_values() {
*coded_spent += codes.combos() - grown.unwrap_or(0);
}
codes.hold(coded_on);
match grown {
Some(span) => coded_map.widen(span, codes.combos()),
None => coded_map.reset(codes.combos()),
}
}
let one_at_a_time = prehashed.is_none() && codes.by_value();
let mut hashed = false;
let mut resolve = |row: usize| -> Result<usize> {
let hash = match prehashed {
Some(prehashed) => prehashed[row],
None if one_at_a_time => codes.hash_of(row),
None => {
if !hashed {
crate::table::hash(
keys,
*length,
hashes,
crate::table::Across::OneInput,
);
hashed = true;
}
hashes[row]
}
};
match table.probe(hash, keys, row) {
Probe::Found(slot) => Ok(slot),
Probe::Vacant(_)
if self.max_groups.is_some_and(|limit| table.len() >= limit) =>
{
Ok(NOWHERE)
}
Probe::Vacant(bucket) => {
let slot = table.insert(bucket, hash, keys, row)?;
*groups = table.len();
self.fresh(states, counts, compact)?;
if self.sets {
self.fresh_seen(seen);
}
Ok(slot)
}
}
};
if codes.place_runs(*length, *length / RUN_ROWS, slot_runs) {
let mut start = 0;
for run in slot_runs.iter_mut() {
let (place, end) = *run;
let mut slot = slot_at(coded_map[place]);
if slot == NOWHERE {
slot = resolve(start)?;
coded_map.set(place, held_at(slot));
}
*run = (slot, end);
start = end;
}
runs_found = true;
} else if codes.look_up(coded_map, nowhere(slots, *length)) != Some(false) {
codes.places(*length, coded_places);
let mut row = 0;
loop {
while row < *length {
let found = slot_at(coded_map[coded_places[row]]);
if found == NOWHERE {
break;
}
slots[row] = found;
row += 1;
}
if row == *length {
break;
}
let slot = resolve(row)?;
slots[row] = slot;
coded_map.set(coded_places[row], held_at(slot));
row += 1;
}
}
} else {
coded_on.clear();
}
if !runs_found && slots.len() != *length {
slots.clear();
slots.resize(*length, if alone { 0 } else { NOWHERE });
}
if !alone && closed.is_none() && direct.is_none() {
match prehashed {
Some(prehashed) => {
hashes.clear();
hashes.extend_from_slice(prehashed);
}
None => crate::table::hash(keys, *length, hashes, crate::table::Across::OneInput),
}
}
if let Some((from, to)) = closed {
crate::table::repeats(keys, *length, 0, same);
let mut slot = NOWHERE;
for row in from..to {
if row == from || !same[row] {
slot = table.append(keys, row)?;
*groups = table.len();
self.fresh(states, counts, compact)?;
}
slots[row] = slot;
}
}
let runs = if closed.is_none()
&& direct.is_none()
&& !alone
&& over.is_none()
&& self.max_groups.is_none()
{
crate::table::repeats(keys, *length, length.div_ceil(2), same)
} else {
same.clear();
0
};
let by_run = runs > 0;
if by_run {
leaders.clear();
leaders.extend((0..*length).filter(|&row| !same[row]));
}
let mut from = 0;
while by_run && from < leaders.len() {
let upto = (from + crate::table::BATCH).min(leaders.len());
let batch = &leaders[from..upto];
from = upto;
leader_slots.clear();
leader_slots.resize(batch.len(), NOWHERE);
table.probe_these(hashes, keys, batch, leader_slots, walk);
for &place in walk.pending() {
let row = batch[place];
let bucket = match table.probe(hashes[row], keys, row) {
Probe::Found(slot) => {
leader_slots[place] = slot;
continue;
}
Probe::Vacant(bucket) => bucket,
};
leader_slots[place] = table.insert(bucket, hashes[row], keys, row)?;
*groups = table.len();
self.fresh(states, counts, compact)?;
if self.sets {
self.fresh_seen(seen);
}
}
for (place, &row) in batch.iter().enumerate() {
slots[row] = leader_slots[place];
}
}
if by_run {
for row in 1..*length {
if same[row] {
slots[row] = slots[row - 1];
}
}
}
let mut from = 0;
while closed.is_none() && !by_run && direct.is_none() && !alone && from < *length {
let upto = (from + crate::table::BATCH).min(*length);
table.probe_run(hashes, keys, from, upto, slots, walk);
from = upto;
for &row in walk.pending() {
let bucket = match table.probe(hashes[row], keys, row) {
Probe::Found(slot) => {
slots[row] = slot;
continue;
}
Probe::Vacant(bucket) => bucket,
};
if self.max_groups.is_some_and(|limit| table.len() >= limit) {
continue;
}
if let Some(file) = over.as_mut() {
put_away(file, seen_rows, row, away)?;
continue;
}
slots[row] = table.insert(bucket, hashes[row], keys, row)?;
*groups = table.len();
self.fresh(states, counts, compact)?;
if self.sets {
self.fresh_seen(seen);
}
}
}
let users = if self.count_only {
1
} else if self.compact_numeric {
0
} else {
self.calls
.iter()
.enumerate()
.filter(|&(at, call)| {
!self.by_vector[at] && call.folds() && !call.distinct && filters[at].is_none()
})
.count()
};
let whole;
let length = match marked {
Some((picks, all)) => {
let mut spread = Vec::with_capacity(slot_runs.len() + 8);
let most = all.saturating_mul(users) / RUN_ROWS;
if runs_found && spread_runs(slot_runs, picks.indices(), *all, most, &mut spread) {
*slot_runs = spread;
} else {
if runs_found {
fill_slots(slots, slot_runs);
}
spread_slots(slots, picks.indices(), *all);
runs_found = false;
}
whole = *all;
&whole
}
None => length,
};
let mut unfilled = runs_found;
let mut fill_now = |slots: &mut Vec<usize>, runs: &[(usize, usize)]| {
if std::mem::take(&mut unfilled) {
fill_slots(slots, runs);
}
};
if self.compact_numeric {
fill_now(slots, slot_runs);
let sum = arguments[1].first().expect("SUM has one argument");
let mean = arguments[2].first().expect("AVG has one argument");
let sum_flat = flat_smallint(sum);
let mean_flat = flat_smallint(mean);
let read =
|column: &Vector, values: Option<&[i16]>, row: usize| -> Result<Option<i16>> {
match values {
Some(values) if column.validity().is_valid(row) => Ok(Some(values[row])),
Some(_) => Ok(None),
None => match column.value_at(row) {
Value::SmallInt(value) => Ok(Some(value)),
Value::Null => Ok(None),
value => Err(Error::internal(format!(
"a compact SMALLINT aggregate received {value:?}"
))),
},
}
};
for (row, &slot) in slots.iter().enumerate() {
if slot == NOWHERE {
continue;
}
let state = &mut compact[slot];
state.add(
slot,
read(sum, sum_flat, row)?,
read(mean, mean_flat, row)?,
overflow,
)?;
}
}
let by_runs = runs_found || slot_runs_of(slots, slot_runs, users);
if self.count_only {
if by_runs {
let mut start = 0;
for &(slot, end) in slot_runs.iter() {
if slot != NOWHERE {
counts[slot] += (end - start) as i64;
}
start = end;
}
} else {
fill_now(slots, slot_runs);
for &slot in slots.iter() {
if slot != NOWHERE {
counts[slot] += 1;
}
}
}
}
let mut counted: Option<Option<Vec<i64>>> = None;
let mut shared = 0_u64;
if by_runs && users > 1 && !self.count_only && !self.compact_numeric {
let offered = self
.calls
.iter()
.enumerate()
.filter(|&(at, call)| {
!self.by_vector[at] && call.folds() && !call.distinct && filters[at].is_none()
})
.fold(0_u64, |offered, (at, _)| offered | one_call(at));
let inputs: Vec<Option<&Vector>> =
arguments.iter().map(|argument| argument.first()).collect();
loop {
let took = update_shared_runs(
states,
slot_runs,
calls,
&inputs,
offered & !shared,
*length,
)?;
if took == 0 {
break;
}
shared |= took;
}
}
for (at, call) in self.calls.iter().enumerate() {
if self.count_only || self.compact_numeric {
break;
}
if self.by_vector[at] || !call.folds() {
continue;
}
if shared & one_call(at) != 0 {
continue;
}
if call.distinct {
fill_now(slots, slot_runs);
aside += self.distinct(states, seen, seen_rows, slots, at, given)?;
continue;
}
if by_runs
&& filters[at].is_none()
&& update_runs(states, slot_runs, calls, at, arguments[at].first(), *length)?
{
continue;
}
fill_now(slots, slot_runs);
let (picked, tallied) = match &filters[at] {
None => {
let rows = slots.len().min(*length);
let groups = states.len().checked_div(calls).unwrap_or(usize::MAX);
let tally = counted.get_or_insert_with(|| group_tally(&slots[..rows], groups));
(&*slots, tally.as_deref())
}
Some(flags) => {
kept.clear();
kept.extend(slots.iter().enumerate().map(|(row, &slot)| {
if slot != NOWHERE && is_true(&flags.value_at(row)) {
slot
} else {
NOWHERE
}
}));
(&*kept, None)
}
};
if !update_general(states, picked, calls, at, &arguments[at], *length)? {
let argument = arguments[at].first();
update_tallied(states, picked, tallied, calls, at, argument, *length)?;
}
}
rows::capacity(table.owned(), charged_keys, scratch)?;
containers.grow(aside)?;
let now = tables(table, states, counts, compact, overflow, seen);
rows::capacity(now, charged, containers)?;
match over.as_ref() {
None if !alone && closed.is_none() && crowded(&self.memory) => {
*over = Some(Spill::new("aggregate", self.spilled_types())?);
}
Some(file) => hopeless(file, *groups)?,
None => {}
}
Ok(())
}
fn finish(
&self,
local: Building,
chunks: &mut Vec<Chunk>,
held: &mut Reservation,
) -> Result<Option<Spill>> {
let timing = stage::Timing::start(Stage::Emit);
let finished = self.finishing(local, chunks, held);
timing.stop(0);
finished
}
fn finishing(
&self,
local: Building,
chunks: &mut Vec<Chunk>,
held: &mut Reservation,
) -> Result<Option<Spill>> {
let Building {
mut scratch,
mut containers,
table,
mut states,
counts,
compact,
overflow,
seen,
groups,
affine_rows,
over,
..
} = local;
let calls = self.calls.len();
drop(seen);
let alive = table.footprint()
+ width_of(states.capacity() * size_of::<Accumulator>())
+ width_of(counts.capacity() * size_of::<i64>())
+ width_of(compact.capacity() * size_of::<CompactNumeric>())
+ overflow_footprint(&overflow);
containers.shrink(containers.bytes().saturating_sub(alive));
let types = self.schema.types();
let width = self.groups.len();
let count = |slot: usize, call: usize| {
if self.count_only {
return counts[slot];
}
if self.compact_numeric {
return compact[slot].count();
}
states[slot * calls + call].counted().expect("a selected COUNT call has a COUNT state")
};
let selected = match (self.top_counts, self.having_count) {
(Some((bound, _)), _) if bound >= groups => None,
(Some((bound, ranks)), _) => {
let mut best = largest(groups, bound, |slot| count(slot, ranks));
best.sort_unstable();
Some(best)
}
(_, Some((call, minimum))) => {
let mut kept = Vec::new();
for slot in 0..groups {
if count(slot, call) >= minimum {
kept.push(slot);
}
}
Some(kept)
}
_ => None,
};
let output_groups = selected.as_ref().map_or(groups, Vec::len);
if !self.count_only && !self.compact_numeric {
settle_extremes(&mut states, selected.as_deref(), groups, calls)?;
}
scratch.grow(width_of(VECTOR_SIZE.min(output_groups) * size_of::<Value>()))?;
let mut results: Vec<Value> = Vec::new();
let mut run: Vec<usize> = Vec::new();
for start in (0..output_groups).step_by(VECTOR_SIZE) {
let end = (start + VECTOR_SIZE).min(output_groups);
let slots = selected.as_ref().map(|slots| &slots[start..end]);
let picked: &[usize] = match slots {
Some(slots) => slots,
None => {
run.clear();
run.extend(start..end);
&run
}
};
let mut columns = Vec::with_capacity(width + calls);
let mut key = 0;
for (at, ty) in types.iter().take(width).enumerate() {
if let Some(value) = &self.constants[at] {
columns.push(Vector::constant(ty.clone(), value.clone(), end - start));
} else {
columns.push(match slots {
Some(slots) => table.column_slots(key, ty, slots)?,
None => table.column(key, ty, start..end)?,
});
key += 1;
}
}
for (at, ty) in types.iter().skip(width).enumerate() {
let held = self.calls[at].state_of(at);
if !self.count_only
&& !self.compact_numeric
&& self.calls[at].finishes_plainly()
&& let Some(vector) = finish_run(&states, picked, calls, held, ty)?
{
columns.push(vector);
continue;
}
let mut taken = 0;
results.clear();
for index in start..end {
let slot = slots.map_or(index, |slots| slots[index - start]);
let value = if self.count_only {
Ok(Value::BigInt(counts[slot]))
} else if self.compact_numeric {
let state = &compact[slot];
let (sum, mean) = state.totals(slot, &overflow);
match at {
0 => Ok(Value::BigInt(state.count())),
1 => Accumulator::exact_sum(
sum,
state.sum_seen(),
&self.calls[1].returns,
)
.finish(),
2 => Accumulator::exact_avg(
mean,
state.mean_count,
&self.calls[2].returns,
)
.finish(),
_ => unreachable!("compact numeric has three calls"),
}
} else {
match (self.calls[at].affine, self.calls[at].reads_total) {
(Some((source, offset)), _) => states[slot * calls + source]
.finish_offset(offset, affine_rows[source]),
(None, Some(source)) => sum_from_mean(
&states[slot * calls + source],
&self.calls[at].returns,
),
(None, None) => states[slot * calls + held].finish(),
}
}?;
taken += rows::owned(&value);
results.push(value);
}
scratch.grow(taken)?;
columns.push(Vector::from_values(ty.clone(), &results)?);
scratch.shrink(taken);
}
let chunk = Chunk::with_rows(columns, end - start)?;
held.grow(width_of(chunk.footprint()))?;
chunks.push(chunk);
}
drop(results);
drop(states);
drop(counts);
drop(compact);
drop(table);
containers.release();
match over {
Some(file) if groups == 0 && file.rows() > 0 => Err(Error::out_of_memory(format!(
"the memory limit does not leave room for a single group of this aggregate, \
{} rows and {} bytes went to a spill file and none of them could be finished",
file.rows(),
file.bytes()
))),
Some(file) if file.rows() > 0 => Ok(Some(file)),
_ => Ok(None),
}
}
fn merge(&self, from: Building, into: &mut Building) -> Result<()> {
let timing = stage::Timing::start(Stage::Merge);
let merged = self.merging(from, into);
timing.stop(0);
merged
}
fn merging(&self, from: Building, into: &mut Building) -> Result<()> {
if from.over.is_some() || into.over.is_some() {
return Err(Error::internal(
"two tables of an aggregate were merged with a spill file between them, where a \
key can be in one table and in the other's file at once, which the callers avoid \
by partitioning instead",
));
}
let Building {
scratch,
containers,
table: source,
states: taken,
counts: tallies,
compact: packed,
overflow: wide,
seen: mut watched,
groups: found,
affine_rows: counted,
..
} = from;
let calls = self.calls.len();
for (at, rows) in counted.iter().enumerate() {
into.affine_rows[at] += rows;
}
let distinct: Vec<bool> = self.calls.iter().map(|call| call.distinct).collect();
let mut coming = Folding {
count_only: self.count_only,
calls,
distinct: &distinct,
taken: &taken,
tallies: &tallies,
compact: &packed,
overflow: &wide,
watched: &mut watched,
};
let mut aside = 0;
if self.alone {
aside += merge_slot(&mut coming, 0, 0, into)?;
} else {
let types: Vec<LogicalType> =
self.keys.iter().map(|&key| self.plan.expr_type(key).clone()).collect();
let mut run: Vec<usize> = Vec::with_capacity(VECTOR_SIZE);
for start in (0..found).step_by(VECTOR_SIZE) {
let end = (start + VECTOR_SIZE).min(found);
let mut keys = Vec::with_capacity(types.len());
for (at, ty) in types.iter().enumerate() {
keys.push(source.column(at, ty, start..end)?);
}
run.clear();
run.extend(start..end);
aside += self.fold_slots(&mut coming, &source, &keys, start, &run, into)?;
}
}
into.containers.grow(aside)?;
drop(watched);
drop(taken);
drop(tallies);
drop(source);
drop(scratch);
drop(containers);
rows::capacity(into.table.owned(), &mut into.charged_keys, &mut into.scratch)?;
let now = tables(
&into.table,
&into.states,
&into.counts,
&into.compact,
&into.overflow,
&into.seen,
);
rows::capacity(now, &mut into.charged, &mut into.containers)
}
fn fold_slots(
&self,
coming: &mut Folding<'_>,
source: &Table,
keys: &[Vector],
start: usize,
slots: &[usize],
into: &mut Building,
) -> Result<u64> {
let mut aside = 0;
let rows: Vec<usize> = slots.iter().map(|&slot| slot - start).collect();
let length = rows.iter().max().map_or(0, |&row| row + 1);
let hashes: Vec<u64> = (start..start + length).map(|slot| source.hash_of(slot)).collect();
let mut targets = vec![usize::MAX; rows.len()];
let mut walk = Walk::default();
for from in (0..rows.len()).step_by(crate::table::BATCH) {
let upto = (from + crate::table::BATCH).min(rows.len());
into.table.probe_these(
&hashes,
keys,
&rows[from..upto],
&mut targets[from..upto],
&mut walk,
);
for &place in walk.pending() {
let row = rows[from + place];
let bucket = match into.table.probe(hashes[row], keys, row) {
Probe::Found(target) => {
targets[from + place] = target;
continue;
}
Probe::Vacant(bucket) => bucket,
};
if self.max_groups.is_some_and(|limit| into.table.len() >= limit) {
continue;
}
targets[from + place] = into.table.insert(bucket, hashes[row], keys, row)?;
into.groups = into.table.len();
self.fresh(&mut into.states, &mut into.counts, &mut into.compact)?;
if self.sets {
self.fresh_seen(&mut into.seen);
}
}
}
for (&slot, &target) in slots.iter().zip(&targets) {
if target != usize::MAX {
aside += merge_slot(coming, slot, target, into)?;
}
}
Ok(aside)
}
fn scatter(&self, from: Building, spin: usize) -> Result<()> {
let timing = stage::Timing::start(Stage::Scatter);
let scattered = self.scattering(from, spin);
timing.stop(0);
scattered
}
fn scattering(&self, from: Building, spin: usize) -> Result<()> {
debug_assert!(from.over.is_none(), "a spill file is drained by hand_over, not scattered");
let Building {
scratch,
containers,
table: source,
states: taken,
counts: tallies,
compact: packed,
overflow: wide,
seen: mut watched,
groups: found,
..
} = from;
let calls = self.calls.len();
let distinct: Vec<bool> = self.calls.iter().map(|call| call.distinct).collect();
let mut coming = Folding {
count_only: self.count_only,
calls,
distinct: &distinct,
taken: &taken,
tallies: &tallies,
compact: &packed,
overflow: &wide,
watched: &mut watched,
};
let types: Vec<LogicalType> =
self.keys.iter().map(|&key| self.plan.expr_type(key).clone()).collect();
let shift = u64::BITS - RADIX_PARTITIONS.ilog2();
let mut buckets: Vec<Vec<usize>> = vec![Vec::new(); RADIX_PARTITIONS];
let mut here: Vec<usize> = Vec::new();
let mut late: Vec<usize> = Vec::new();
for start in (0..found).step_by(VECTOR_SIZE) {
let end = (start + VECTOR_SIZE).min(found);
let mut keys = Vec::with_capacity(types.len());
for (at, ty) in types.iter().enumerate() {
keys.push(source.column(at, ty, start..end)?);
}
for bucket in &mut buckets {
bucket.clear();
}
for slot in start..end {
buckets[(source.hash_of(slot) >> shift) as usize].push(slot);
}
for step in 0..RADIX_PARTITIONS {
let at = (step + spin) % RADIX_PARTITIONS;
if buckets[at].is_empty() {
continue;
}
let mut held = self.merged[at].lock().map_err(poisoned)?;
let Partition { table, carried, .. } = &mut *held;
let into = table.get_or_insert_with(|| self.partition());
if into.over.is_none() {
let grown =
self.fold_slots(&mut coming, &source, &keys, start, &buckets[at], into)?;
charge(into, grown)?;
continue;
}
here.clear();
late.clear();
for &slot in &buckets[at] {
match into.table.probe(source.hash_of(slot), &keys, slot - start) {
Probe::Found(_) => here.push(slot),
Probe::Vacant(_) => late.push(slot),
}
}
let grown = self.fold_slots(&mut coming, &source, &keys, start, &here, into)?;
charge(into, grown)?;
let waiting = carried.get_or_insert_with(|| self.partition());
let grown = self.fold_slots(&mut coming, &source, &keys, start, &late, waiting)?;
charge(waiting, grown)?;
}
}
drop(watched);
drop(taken);
drop(tallies);
drop(source);
drop(scratch);
drop(containers);
Ok(())
}
fn agree(
&self,
rows: &Rows,
limit: usize,
into: &mut Building,
installed: &mut bool,
) -> Result<()> {
if !self.settled.load(Ordering::Acquire) {
self.collect(rows, limit)?;
}
if *installed || !self.settled.load(Ordering::Acquire) {
return Ok(());
}
let keys = {
let held = self.agreed.lock().map_err(poisoned)?;
let agreed = held
.as_ref()
.ok_or_else(|| Error::internal("a limited aggregate settled on nothing"))?;
let keys = agreed
.keys
.as_ref()
.ok_or_else(|| Error::internal("a limited aggregate settled without keys"))?;
Arc::clone(keys)
};
self.install(&keys, into)?;
*installed = true;
Ok(())
}
fn collect(&self, rows: &Rows, limit: usize) -> Result<()> {
let mut held = self.agreed.lock().map_err(poisoned)?;
let agreed = match held.as_mut() {
Some(agreed) => agreed,
None => {
let types: Vec<LogicalType> =
rows.keys.iter().map(|column| column.logical_type().clone()).collect();
held.insert(Agreed { table: Table::new(&types), hashes: Vec::new(), keys: None })
}
};
if agreed.keys.is_some() {
return Ok(());
}
crate::table::hash(
&rows.keys,
rows.rows,
&mut agreed.hashes,
crate::table::Across::OneInput,
);
for row in 0..rows.rows {
if agreed.table.len() >= limit {
break;
}
let hash = agreed.hashes[row];
if let Probe::Vacant(bucket) = agreed.table.probe(hash, &rows.keys, row) {
agreed.table.insert(bucket, hash, &rows.keys, row)?;
}
}
if agreed.table.len() < limit {
return Ok(());
}
let mut keys = Vec::with_capacity(rows.keys.len());
for (at, column) in rows.keys.iter().enumerate() {
keys.push(agreed.table.column(at, column.logical_type(), 0..agreed.table.len())?);
}
agreed.keys = Some(Arc::new(keys));
self.settled.store(true, Ordering::Release);
Ok(())
}
fn install(&self, keys: &[Vector], into: &mut Building) -> Result<()> {
let rows = keys.first().map_or(0, Vector::len);
let Building { table, states, counts, compact, seen, groups, hashes, .. } = into;
crate::table::hash(keys, rows, hashes, crate::table::Across::OneInput);
for (row, &hash) in hashes.iter().enumerate().take(rows) {
if let Probe::Vacant(bucket) = table.probe(hash, keys, row) {
table.insert(bucket, hash, keys, row)?;
*groups = table.len();
self.fresh(states, counts, compact)?;
if self.sets {
self.fresh_seen(seen);
}
}
}
Ok(())
}
fn ought_to_partition(&self, table: &Building) -> bool {
!self.alone
&& self.max_groups.is_none()
&& self.started.load(Ordering::Relaxed) > 1
&& (table.groups >= self.partition_from || crowded(&self.memory))
}
fn begin_partitioning(
&self,
spreading: &mut Spreading,
own: &mut [Option<Building>],
) -> Result<()> {
let mut built = self.built.lock().map_err(poisoned)?;
if built.partitioning {
return Ok(());
}
built.partitioning = true;
let seeded = self.merged[0].lock().map_err(poisoned)?.table.take();
drop(built);
if !self.worth_local() {
self.give_up_local(spreading)?;
}
match seeded {
Some(seeded) => self.hand(seeded, spreading, own),
None => Ok(()),
}
}
fn hand(
&self,
from: Building,
spreading: &mut Spreading,
own: &mut [Option<Building>],
) -> Result<()> {
if from.over.is_some() || crowded(&self.memory) {
self.give_up_local(spreading)?;
self.hand_all(spreading, own)?;
}
if self.locally.load(Ordering::Relaxed) {
return self.scatter_own(from, own);
}
self.hand_over(from, spreading)
}
fn hand_all(&self, spreading: &mut Spreading, own: &mut [Option<Building>]) -> Result<()> {
for held in own.iter_mut() {
let Some(table) = held.take() else { continue };
self.hand_over(table, spreading)?;
}
Ok(())
}
fn still_local(
&self,
folded: u64,
spreading: &mut Spreading,
own: &mut [Option<Building>],
) -> Result<bool> {
let spilled = own.iter().flatten().any(|table| table.over.is_some());
let mine: u64 = own
.iter()
.flatten()
.map(|table| table.scratch.bytes() + table.containers.bytes())
.sum();
let groups: u64 = own.iter().flatten().map(|table| table.groups as u64).sum();
if !spilled
&& !crowded(&self.memory)
&& self.room_for_local(mine)
&& (self.cache_holds_local(mine)
|| !copies_overlap(folded, groups)
|| keys_arrive_together(spreading))
{
return Ok(true);
}
self.give_up_local(spreading)?;
self.hand_all(spreading, own)?;
Ok(false)
}
fn worth_local(&self) -> bool {
if self.keys.iter().any(|&key| self.plan.expr_type(key) == &LogicalType::Varchar) {
return false;
}
let Some(limit) = self.memory.limit() else { return true };
let instances = self.started.load(Ordering::Relaxed).max(self.merged.len()) as u64;
self.memory.used().saturating_mul(instances) < limit / 2
}
fn room_for_local(&self, mine: u64) -> bool {
let Some(limit) = self.memory.limit() else { return true };
let instances = self.started.load(Ordering::Relaxed) as u64;
mine.saturating_mul(instances) < limit / 4
}
fn cache_holds_local(&self, mine: u64) -> bool {
let instances = self.started.load(Ordering::Relaxed) as u64;
mine.saturating_mul(instances) <= LOCAL_CACHE
}
fn give_up_local(&self, spreading: &mut Spreading) -> Result<()> {
let mut built = self.built.lock().map_err(poisoned)?;
if !built.local {
return Ok(());
}
built.local = false;
self.locally.store(false, Ordering::Relaxed);
let mut handed = Vec::new();
for partition in &self.merged {
handed.append(&mut partition.lock().map_err(poisoned)?.pending);
}
drop(built);
for table in handed {
self.hand_over(table, spreading)?;
}
Ok(())
}
fn deposit(&self, own: &mut [Option<Building>], spreading: &mut Spreading) -> Result<()> {
if own.iter().flatten().next().is_none() {
return Ok(());
}
let built = self.built.lock().map_err(poisoned)?;
if built.local {
for (at, held) in own.iter_mut().enumerate() {
let Some(table) = held.take() else { continue };
self.merged[at].lock().map_err(poisoned)?.pending.push(table);
}
return Ok(());
}
drop(built);
self.hand_all(spreading, own)
}
fn hand_over(&self, mut from: Building, spreading: &mut Spreading) -> Result<()> {
let leftover = from.over.take();
self.scatter(from, spreading.spin)?;
let Some(mut file) = leftover else { return Ok(()) };
let mut spilled = Spilled::new(file.read()?, self.spilled_types());
while let Some(rows) = spilled.next(self)? {
self.spread(&rows, spreading)?;
}
Ok(())
}
fn drain(
&self,
folded: u64,
spreading: &mut Spreading,
own: &mut [Option<Building>],
) -> Result<()> {
if spreading.pending.is_empty() {
return Ok(());
}
let pieces = std::mem::take(&mut spreading.pending);
spreading.gathered = 0;
for rows in Rows::joined(pieces)? {
if self.locally.load(Ordering::Relaxed) && self.still_local(folded, spreading, own)? {
self.spread_own(&rows, spreading, own)?;
} else {
self.spread(&rows, spreading)?;
}
}
Ok(())
}
fn split(&self, rows: &Rows, spreading: &mut Spreading) -> Result<Vec<Option<Rows>>> {
let Spreading { hashes, picks, keyed, spin, split_rows, runs, .. } = spreading;
crate::table::hash(&rows.keys, rows.rows, hashes, crate::table::Across::OneInput);
for pick in picks.iter_mut() {
pick.clear();
}
for hashed in keyed.iter_mut() {
hashed.clear();
}
let shift = u64::BITS - RADIX_PARTITIONS.ilog2();
let mut before: Option<u64> = None;
for (row, &hash) in hashes.iter().enumerate() {
if before != Some(hash) {
*runs += 1;
before = Some(hash);
}
let partition = (hash >> shift) as usize;
picks[partition].push(row as u32);
keyed[partition].push(hash);
}
*split_rows += rows.rows as u64;
let mut ready: Vec<Option<Rows>> = Vec::with_capacity(RADIX_PARTITIONS);
for pick in picks.iter() {
ready.push(if pick.is_empty() { None } else { Some(rows.gather(pick)?) });
}
*spin = (*spin + 1) % RADIX_PARTITIONS;
Ok(ready)
}
fn spread_own(
&self,
rows: &Rows,
spreading: &mut Spreading,
own: &mut [Option<Building>],
) -> Result<()> {
let timing = stage::Timing::start(Stage::Fold);
let spread = self.spreading_own(rows, spreading, own);
timing.stop(0);
spread
}
fn spreading_own(
&self,
rows: &Rows,
spreading: &mut Spreading,
own: &mut [Option<Building>],
) -> Result<()> {
let ready = self.split(rows, spreading)?;
for (partition, selected) in ready.iter().enumerate() {
let Some(selected) = selected else { continue };
let table = own[partition].get_or_insert_with(|| self.kept());
if let Some(error) = table.failure.take() {
return Err(error);
}
self.fold(selected, table, Some(&spreading.keyed[partition]), None)?;
}
Ok(())
}
fn scatter_own(&self, from: Building, own: &mut [Option<Building>]) -> Result<()> {
let timing = stage::Timing::start(Stage::Scatter);
let scattered = self.scattering_own(from, own);
timing.stop(0);
scattered
}
fn scattering_own(&self, from: Building, own: &mut [Option<Building>]) -> Result<()> {
debug_assert!(from.over.is_none(), "a spilled table is never scattered locally");
let Building {
scratch,
containers,
table: source,
states: taken,
counts: tallies,
compact: packed,
overflow: wide,
seen: mut watched,
groups: found,
..
} = from;
let calls = self.calls.len();
let distinct: Vec<bool> = self.calls.iter().map(|call| call.distinct).collect();
let mut coming = Folding {
count_only: self.count_only,
calls,
distinct: &distinct,
taken: &taken,
tallies: &tallies,
compact: &packed,
overflow: &wide,
watched: &mut watched,
};
let types: Vec<LogicalType> =
self.keys.iter().map(|&key| self.plan.expr_type(key).clone()).collect();
let shift = u64::BITS - RADIX_PARTITIONS.ilog2();
let mut buckets: Vec<Vec<usize>> = vec![Vec::new(); RADIX_PARTITIONS];
for start in (0..found).step_by(VECTOR_SIZE) {
let end = (start + VECTOR_SIZE).min(found);
let mut keys = Vec::with_capacity(types.len());
for (at, ty) in types.iter().enumerate() {
keys.push(source.column(at, ty, start..end)?);
}
for bucket in &mut buckets {
bucket.clear();
}
for slot in start..end {
buckets[(source.hash_of(slot) >> shift) as usize].push(slot);
}
for (at, bucket) in buckets.iter().enumerate() {
if bucket.is_empty() {
continue;
}
let into = own[at].get_or_insert_with(|| self.kept());
let grown = self.fold_slots(&mut coming, &source, &keys, start, bucket, into)?;
charge(into, grown)?;
}
}
drop(watched);
drop(taken);
drop(tallies);
drop(source);
drop(scratch);
drop(containers);
Ok(())
}
fn spread(&self, rows: &Rows, spreading: &mut Spreading) -> Result<()> {
let timing = stage::Timing::start(Stage::Fold);
let spread = self.spreading(rows, spreading);
timing.stop(0);
spread
}
fn spreading(&self, rows: &Rows, spreading: &mut Spreading) -> Result<()> {
let ready = self.split(rows, spreading)?;
let Spreading { keyed, spin, waiting, .. } = spreading;
let spin = &*spin;
waiting.clear();
for step in 0..RADIX_PARTITIONS {
let partition = (step + *spin) % RADIX_PARTITIONS;
let Some(selected) = &ready[partition] else { continue };
match self.merged[partition].try_lock() {
Ok(mut held) => {
let table = held.table.get_or_insert_with(|| self.partition());
self.fold(selected, table, Some(&keyed[partition]), None)?;
}
Err(TryLockError::WouldBlock) => waiting.push(partition),
Err(TryLockError::Poisoned(error)) => return Err(poisoned(error)),
}
}
for &partition in waiting.iter() {
let selected =
ready[partition].as_ref().expect("only a filled partition was put aside");
let mut held = self.merged[partition].lock().map_err(poisoned)?;
let table = held.table.get_or_insert_with(|| self.partition());
self.fold(selected, table, Some(&keyed[partition]), None)?;
}
Ok(())
}
fn distinct(
&self,
states: &mut [Accumulator],
seen: &mut [DistinctSet],
rows: &Rows,
slots: &[usize],
at: usize,
given: &mut [Key],
) -> Result<u64> {
let calls = self.calls.len();
let mut aside = 0;
for (row, &slot) in slots.iter().enumerate() {
if slot == NOWHERE {
continue;
}
if let Some(flags) = &rows.filters[at]
&& !is_true(&flags.value_at(row))
{
continue;
}
if let (DistinctSet::BigInt(set), [column]) =
(&mut seen[slot * calls + at], rows.arguments[at].as_slice())
{
if column.is_null_at(row) {
continue;
}
let value = match column.signed_at(row) {
Some(value) => i64::try_from(value)
.map_err(|_| Error::internal("a BIGINT distinct value is out of range"))?,
None => match column.try_value_at(row)? {
Value::BigInt(value) => value,
value => {
return Err(Error::internal(format!(
"a BIGINT distinct set was given {value:?}"
)));
}
},
};
if set.insert(value) {
aside += width_of(size_of::<i64>() * 2);
states[slot * calls + at].update(&[Value::BigInt(value)])?;
}
continue;
}
let args = &mut given[at];
fill(args, &rows.arguments[at], row)?;
let DistinctSet::Row(set) = &mut seen[slot * calls + at] else {
return Err(Error::internal("a distinct set did not match its argument"));
};
if set.contains(args) {
continue;
}
let stored = args.clone();
aside += rows::footprint(&stored.0);
set.insert(stored);
states[slot * calls + at].update(&args.0)?;
}
Ok(aside)
}
fn fresh(
&self,
states: &mut Vec<Accumulator>,
counts: &mut Vec<i64>,
compact: &mut Vec<CompactNumeric>,
) -> Result<()> {
if self.count_only {
counts.push(0);
return Ok(());
}
if self.compact_numeric {
compact.push(CompactNumeric::default());
return Ok(());
}
if let Some(template) = &self.template {
states.extend_from_slice(template);
return Ok(());
}
for call in &self.calls {
states.push(Accumulator::new(&call.name, &call.returns)?);
}
Ok(())
}
fn fresh_seen(&self, seen: &mut Vec<DistinctSet>) {
for call in &self.calls {
let big_int = call.distinct
&& call.args.len() == 1
&& self.plan.expr_type(call.args[0]) == &LogicalType::BigInt;
if big_int {
seen.push(DistinctSet::BigInt(BigIntDistinct::default()));
} else {
seen.push(DistinctSet::Row(RowSet::default()));
}
}
}
fn close_runs(&self, rows: &Rows, from: usize, to: usize) -> Result<Option<Chunk>> {
let mut same = Vec::new();
crate::table::repeats(&rows.keys, rows.rows, 0, &mut same);
let mut starts: Vec<u32> = Vec::new();
for (row, &repeat) in same.iter().enumerate().take(to).skip(from) {
if row == from || !repeat {
starts.push(u32::try_from(row).map_err(|_| Error::internal("a chunk too long"))?);
}
}
let groups = starts.len();
let end = |group: usize| starts.get(group + 1).map_or(to, |&start| start as usize);
let types = self.schema.types();
let width = self.groups.len();
let mut columns = Vec::with_capacity(types.len());
let mut key = 0;
for (at, ty) in types.iter().take(width).enumerate() {
if let Some(value) = &self.constants[at] {
columns.push(Vector::constant(ty.clone(), value.clone(), groups));
} else {
columns.push(rows.keys[key].gather(&starts)?);
key += 1;
}
}
let mut answers = vec![0_i128; groups];
let mut valid = vec![true; groups];
for (at, ty) in types.iter().skip(width).enumerate() {
valid.fill(true);
if self.calls[at].name == "count_star" {
for (group, &start) in starts.iter().enumerate() {
answers[group] = (end(group) - start as usize) as i128;
}
} else {
let argument = rows.arguments[at][0].clone().into_flat()?;
let nulls = argument.validity().has_nulls(rows.rows).then(|| argument.validity());
let counting = self.calls[at].name == "count";
let values = match counting {
true => None,
false => match integers(&argument, rows.rows) {
Some(values) => Some(values),
None => return Ok(None),
},
};
for (group, &start) in starts.iter().enumerate() {
let run = start as usize..end(group);
let (total, seen) = match (&values, nulls) {
(None, None) => (run.len() as i128, true),
(None, Some(nulls)) => {
(run.filter(|&row| nulls.is_valid(row)).count() as i128, true)
}
(Some(values), None) => {
(values[run].iter().map(|&v| i128::from(v)).sum(), true)
}
(Some(values), Some(nulls)) => {
let mut total = 0_i128;
let mut seen = false;
for row in run {
if nulls.is_valid(row) {
total += i128::from(values[row]);
seen = true;
}
}
(total, seen)
}
};
answers[group] = total;
valid[group] = seen;
}
}
columns.push(whole_answers(&answers, &valid, ty)?);
}
Chunk::with_rows(columns, groups).map(Some)
}
fn spilled_types(&self) -> Vec<LogicalType> {
let mut types = Vec::new();
for &group in &self.keys {
types.push(self.plan.expr_type(group).clone());
}
for call in &self.calls {
for &argument in &call.args {
types.push(self.plan.expr_type(argument).clone());
}
}
for call in &self.calls {
if let Some(filter) = call.filter {
types.push(self.plan.expr_type(filter).clone());
}
}
types
}
}
fn nowhere(slots: &mut Vec<usize>, rows: usize) -> &mut [usize] {
slots.clear();
slots.resize(rows, NOWHERE);
slots
}
fn fill_slots(slots: &mut Vec<usize>, runs: &[(usize, usize)]) {
slots.clear();
for &(slot, end) in runs {
slots.resize(end, slot);
}
}
fn spread_slots(slots: &mut Vec<usize>, kept: &[u32], all: usize) {
slots.resize(all, NOWHERE);
for (place, &row) in kept.iter().enumerate().rev() {
let slot = std::mem::replace(&mut slots[place], NOWHERE);
slots[row as usize] = slot;
}
}
fn spread_runs(
runs: &[(usize, usize)],
kept: &[u32],
all: usize,
most: usize,
into: &mut Vec<(usize, usize)>,
) -> bool {
into.clear();
let mut push = |slot: usize, end: usize| {
match into.last_mut() {
Some(last) if last.0 == slot => last.1 = end,
_ => into.push((slot, end)),
}
into.len() <= most
};
let mut from = 0;
let mut row = 0;
for &(slot, end) in runs {
let Some(within) = kept.get(from..end) else { return false };
let mut rest = within;
while let (Some(&first), Some(&last)) = (rest.first(), rest.last()) {
let (first, last) = (first as usize, last as usize);
let length = if last - first == rest.len() - 1 {
rest.len()
} else {
1 + rest.windows(2).take_while(|pair| pair[1] == pair[0] + 1).count()
};
let upto = first + length;
if (first > row && !push(NOWHERE, first)) || !push(slot, upto) {
into.clear();
return false;
}
row = upto;
rest = &rest[length..];
}
from = end;
}
if from != kept.len() || row > all || (row < all && !push(NOWHERE, all)) {
into.clear();
return false;
}
true
}
#[derive(Clone, Debug)]
struct Rows {
keys: Vec<Vector>,
arguments: Vec<Vec<Vector>>,
filters: Vec<Option<Vector>>,
rows: usize,
marked: Option<(Selection, usize)>,
}
impl Rows {
fn settled(&self) -> Result<Cow<'_, Self>> {
let Some((kept, _)) = &self.marked else { return Ok(Cow::Borrowed(self)) };
let cut = |vector: &Vector| vector.gather(kept.indices());
Ok(Cow::Owned(Self {
keys: self.keys.clone(),
arguments: self
.arguments
.iter()
.map(|call| call.iter().map(cut).collect::<Result<_>>())
.collect::<Result<_>>()?,
filters: self
.filters
.iter()
.map(|filter| filter.as_ref().map(cut).transpose())
.collect::<Result<_>>()?,
rows: self.rows,
marked: None,
}))
}
fn slice(&self, at: usize, len: usize) -> Result<Self> {
let cut = |vector: &Vector| vector.slice(at, len);
Ok(Self {
keys: self.keys.iter().map(cut).collect::<Result<_>>()?,
arguments: self
.arguments
.iter()
.map(|call| call.iter().map(cut).collect::<Result<_>>())
.collect::<Result<_>>()?,
filters: self
.filters
.iter()
.map(|filter| filter.as_ref().map(cut).transpose())
.collect::<Result<_>>()?,
rows: len,
marked: None,
})
}
fn width(&self) -> usize {
self.keys.len()
+ self.arguments.iter().map(Vec::len).sum::<usize>()
+ self.filters.iter().flatten().count()
}
fn joined(pieces: Vec<Self>) -> Result<Vec<Self>> {
if pieces.len() < 2 {
return Ok(pieces);
}
let first = &pieces[0];
let mut keys = Vec::with_capacity(first.keys.len());
for at in 0..first.keys.len() {
let Some(laid) = lay(&pieces, |piece| piece.keys.get(at))? else { return Ok(pieces) };
keys.push(laid);
}
let mut arguments = Vec::with_capacity(first.arguments.len());
for (call, columns) in first.arguments.iter().enumerate() {
let mut laid_call = Vec::with_capacity(columns.len());
for at in 0..columns.len() {
let column = |piece| Self::argument(piece, call, at);
let Some(laid) = lay(&pieces, column)? else { return Ok(pieces) };
laid_call.push(laid);
}
arguments.push(laid_call);
}
let mut filters = Vec::with_capacity(first.filters.len());
for (call, filter) in first.filters.iter().enumerate() {
if filter.is_none() {
if pieces.iter().any(|piece| piece.filters.get(call).is_none_or(Option::is_some)) {
return Ok(pieces);
}
filters.push(None);
continue;
}
let column = |piece| Self::filter(piece, call);
let Some(laid) = lay(&pieces, column)? else { return Ok(pieces) };
filters.push(Some(laid));
}
let rows = pieces.iter().map(|piece| piece.rows).sum();
Ok(vec![Self { keys, arguments, filters, rows, marked: None }])
}
fn argument(&self, call: usize, at: usize) -> Option<&Vector> {
self.arguments.get(call)?.get(at)
}
fn filter(&self, call: usize) -> Option<&Vector> {
self.filters.get(call)?.as_ref()
}
fn gather(&self, rows: &[u32]) -> Result<Self> {
Ok(Self {
keys: self.keys.iter().map(|column| column.gather(rows)).collect::<Result<_>>()?,
arguments: self
.arguments
.iter()
.map(|arguments| {
arguments.iter().map(|column| column.gather(rows)).collect::<Result<_>>()
})
.collect::<Result<_>>()?,
filters: self
.filters
.iter()
.map(|filter| filter.as_ref().map(|column| column.gather(rows)).transpose())
.collect::<Result<_>>()?,
rows: rows.len(),
marked: None,
})
}
}
fn lay<'a>(
pieces: &'a [Rows],
column: impl Fn(&'a Rows) -> Option<&'a Vector>,
) -> Result<Option<Vector>> {
let mut columns = Vec::with_capacity(pieces.len());
for piece in pieces {
let Some(vector) = column(piece) else { return Ok(None) };
columns.push(vector);
}
rudb_vector::concat(columns[0].logical_type(), &columns)
}
#[derive(Debug)]
struct Agreed {
table: Table,
hashes: Vec<u64>,
keys: Option<Arc<Vec<Vector>>>,
}
#[derive(Debug)]
pub(crate) struct Partitioned {
mixed: group_mixed::Local,
counted: group_count::Local,
ranged: group_ranged::Local,
grouped_distinct: group_distinct::Local,
encoded: bool,
encoded_records: EncodedRecords,
encoded_memory: Reservation,
radix_distinct: bool,
radix_distinct_records: Vec<BigIntDistinctPartition>,
radix_distinct_memory: Reservation,
fixed: bool,
fixed_records: Vec<FixedRun>,
fixed_memory: Reservation,
fixed_blocks: FixedBlocks,
dense: bool,
dense_codes: Vec<Blocks<u32>>,
dense_nulls: i64,
dense_memory: Reservation,
single: Option<Building>,
closed: Option<Building>,
ran: Vec<Chunk>,
ran_memory: Reservation,
installed: bool,
expressions: Scratch,
spreading: Spreading,
own: Vec<Option<Building>>,
folded: u64,
}
#[derive(Debug, Default)]
struct FixedBlocks {
held: [SignedBlock; 4],
}
impl FixedBlocks {
fn read(&mut self, rows: usize, columns: [&Vector; 4]) -> Result<()> {
for (held, column) in self.held.iter_mut().zip(columns) {
held.read(rows, column)?;
}
Ok(())
}
fn nulled(&self) -> [bool; 4] {
let [first, second, sum, mean] = &self.held;
[first.nulled(), second.nulled(), sum.nulled(), mean.nulled()]
}
fn cut(&self, rows: usize) -> Result<[&[i64]; 4]> {
let [first, second, sum, mean] = &self.held;
Ok([first.cut(rows)?, second.cut(rows)?, sum.cut(rows)?, mean.cut(rows)?])
}
}
#[derive(Debug)]
struct Spreading {
hashes: Vec<u64>,
picks: Vec<Vec<u32>>,
keyed: Vec<Vec<u64>>,
spin: usize,
waiting: Vec<usize>,
split_rows: u64,
runs: u64,
pending: Vec<Rows>,
gathered: usize,
}
impl Spreading {
fn new() -> Self {
Self {
hashes: Vec::new(),
picks: vec![Vec::new(); RADIX_PARTITIONS],
keyed: vec![Vec::new(); RADIX_PARTITIONS],
spin: 0,
waiting: Vec::new(),
split_rows: 0,
runs: 0,
pending: Vec::new(),
gathered: 0,
}
}
}
fn whole_total(argument: &LogicalType, returns: &LogicalType) -> bool {
use LogicalType as T;
match (argument, returns) {
(T::Decimal { width, scale }, T::Decimal { scale: declared, .. }) => {
*width <= 18 && scale == declared
}
(
T::TinyInt
| T::SmallInt
| T::Integer
| T::BigInt
| T::UTinyInt
| T::USmallInt
| T::UInteger,
T::TinyInt
| T::SmallInt
| T::Integer
| T::BigInt
| T::HugeInt
| T::UTinyInt
| T::USmallInt
| T::UInteger
| T::UBigInt
| T::UHugeInt,
) => true,
_ => false,
}
}
fn integers(flat: &Vector, rows: usize) -> Option<Vec<i64>> {
macro_rules! widened {
($values:expr) => {{
let values = $values.as_slice();
values.get(..rows)?.iter().map(|&value| i64::from(value)).collect()
}};
}
Some(match flat.data() {
Some(Data::Int8(values)) => widened!(values),
Some(Data::Int16(values)) => widened!(values),
Some(Data::Int32(values)) => widened!(values),
Some(Data::Int64(values)) => widened!(values),
Some(Data::UInt8(values)) => widened!(values),
Some(Data::UInt16(values)) => widened!(values),
Some(Data::UInt32(values)) => widened!(values),
_ => return None,
})
}
fn interior(key: &Vector, rows: usize, grouped: bool) -> Option<(usize, usize)> {
if rows < 3 || key.validity().has_nulls(rows) {
return None;
}
fn bounds<T: PartialOrd>(
rows: usize,
grouped: bool,
at: impl Fn(usize) -> T,
) -> Option<(usize, usize)> {
let (mut from, mut to) = (0, 0);
let mut before = at(0);
for row in 1..rows {
let value = at(row);
if value < before && !grouped {
return None;
}
if value != before {
if from == 0 {
from = row;
}
to = row;
}
before = value;
}
(from > 0 && from < to).then_some((from, to))
}
macro_rules! flat {
($values:expr) => {{
let values = $values.as_slice();
if values.len() < rows {
return None;
}
bounds(rows, grouped, |row| values[row])
}};
}
if let Some(packed) = key.packed_parts() {
return bounds(rows, grouped, |row| packed.code(row));
}
match key.data()? {
Data::Int8(values) => flat!(values),
Data::Int16(values) => flat!(values),
Data::Int32(values) => flat!(values),
Data::Int64(values) => flat!(values),
Data::UInt8(values) => flat!(values),
Data::UInt16(values) => flat!(values),
Data::UInt32(values) => flat!(values),
Data::UInt64(values) => flat!(values),
_ => None,
}
}
#[derive(Debug)]
pub(crate) struct Building {
scratch: Reservation,
containers: Reservation,
charged: u64,
charged_keys: u64,
table: Table,
states: Vec<Accumulator>,
counts: Vec<i64>,
compact: Vec<CompactNumeric>,
overflow: HashMap<usize, (i128, i128)>,
seen: Vec<DistinctSet>,
groups: usize,
given: Vec<Key>,
hashes: Vec<u64>,
slots: Vec<usize>,
walk: Walk,
coded_on: Vec<Origin>,
coded_map: Places,
coded_places: Vec<usize>,
coded_values: crate::table::Widened,
slot_runs: Vec<(usize, usize)>,
coded_spent: usize,
coded_read: usize,
same: Vec<bool>,
leaders: Vec<usize>,
leader_slots: Vec<usize>,
kept: Vec<usize>,
affine_rows: Vec<i64>,
over: Option<Spill>,
away: Vec<Value>,
failure: Option<Error>,
}
#[derive(Debug, Default, Clone)]
struct CompactNumeric {
count_and_sum_seen: u64,
sum: i64,
mean: i64,
mean_count: i64,
}
impl CompactNumeric {
const SUM_SEEN: u64 = 1 << 63;
const COUNT: u64 = Self::SUM_SEEN - 1;
fn count(&self) -> i64 {
(self.count_and_sum_seen & Self::COUNT) as i64
}
fn sum_seen(&self) -> bool {
self.count_and_sum_seen & Self::SUM_SEEN != 0
}
fn totals(&self, slot: usize, overflow: &HashMap<usize, (i128, i128)>) -> (i128, i128) {
if overflow.is_empty() {
return (i128::from(self.sum), i128::from(self.mean));
}
overflow.get(&slot).copied().unwrap_or((i128::from(self.sum), i128::from(self.mean)))
}
fn add(
&mut self,
slot: usize,
sum: Option<i16>,
mean: Option<i16>,
overflow: &mut HashMap<usize, (i128, i128)>,
) -> Result<()> {
let count = self
.count()
.checked_add(1)
.ok_or_else(|| Error::out_of_range("a compact COUNT overflowed BIGINT"))?;
self.count_and_sum_seen =
count as u64 | if self.sum_seen() || sum.is_some() { Self::SUM_SEEN } else { 0 };
self.mean_count = self
.mean_count
.checked_add(i64::from(mean.is_some()))
.ok_or_else(|| Error::out_of_range("a compact AVG count overflowed BIGINT"))?;
let added_sum = i64::from(sum.unwrap_or(0));
let added_mean = i64::from(mean.unwrap_or(0));
if overflow.is_empty()
&& let (Some(total_sum), Some(total_mean)) =
(self.sum.checked_add(added_sum), self.mean.checked_add(added_mean))
{
self.sum = total_sum;
self.mean = total_mean;
return Ok(());
}
if let std::collections::hash_map::Entry::Vacant(entry) = overflow.entry(slot) {
if let (Some(total_sum), Some(total_mean)) =
(self.sum.checked_add(added_sum), self.mean.checked_add(added_mean))
{
self.sum = total_sum;
self.mean = total_mean;
return Ok(());
}
entry.insert((
i128::from(self.sum) + i128::from(added_sum),
i128::from(self.mean) + i128::from(added_mean),
));
return Ok(());
}
let totals = overflow.get_mut(&slot).expect("a wide compact total has an overflow entry");
totals.0 = totals
.0
.checked_add(i128::from(added_sum))
.ok_or_else(|| Error::out_of_range("a compact SUM overflowed its exact total"))?;
totals.1 = totals
.1
.checked_add(i128::from(added_mean))
.ok_or_else(|| Error::out_of_range("a compact AVG overflowed its exact total"))?;
Ok(())
}
fn combine(
&mut self,
target: usize,
coming: &Self,
slot: usize,
from: &HashMap<usize, (i128, i128)>,
into: &mut HashMap<usize, (i128, i128)>,
) -> Result<()> {
let (sum, mean) = self.totals(target, into);
let (coming_sum, coming_mean) = coming.totals(slot, from);
let sum = sum
.checked_add(coming_sum)
.ok_or_else(|| Error::out_of_range("a compact SUM overflowed its exact total"))?;
let mean = mean
.checked_add(coming_mean)
.ok_or_else(|| Error::out_of_range("a compact AVG overflowed its exact total"))?;
let count = self
.count()
.checked_add(coming.count())
.ok_or_else(|| Error::out_of_range("a compact COUNT overflowed BIGINT"))?;
self.count_and_sum_seen =
count as u64 | if self.sum_seen() || coming.sum_seen() { Self::SUM_SEEN } else { 0 };
self.mean_count = self
.mean_count
.checked_add(coming.mean_count)
.ok_or_else(|| Error::out_of_range("a compact AVG count overflowed BIGINT"))?;
match (i64::try_from(sum), i64::try_from(mean)) {
(Ok(sum), Ok(mean)) => {
self.sum = sum;
self.mean = mean;
into.remove(&target);
}
_ => {
into.insert(target, (sum, mean));
}
}
Ok(())
}
}
fn flat_smallint(column: &Vector) -> Option<&[i16]> {
match column.data() {
Some(Data::Int16(values)) => Some(values.as_slice()),
_ => None,
}
}
struct Spilled<'s> {
reader: Reader<'s>,
types: Vec<LogicalType>,
row: Vec<Value>,
columns: Vec<Vec<Value>>,
}
#[derive(Debug)]
enum DistinctSet {
BigInt(BigIntDistinct),
Row(RowSet),
}
#[derive(Debug, Default)]
enum BigIntDistinct {
#[default]
Empty,
One(i64),
Many(BigIntSet),
}
impl BigIntDistinct {
fn insert(&mut self, value: i64) -> bool {
match self {
Self::Empty => {
*self = Self::One(value);
true
}
Self::One(held) if *held == value => false,
Self::One(held) => {
let first = *held;
let mut values = BigIntSet::default();
values.insert(first);
values.insert(value);
*self = Self::Many(values);
true
}
Self::Many(values) => values.insert(value),
}
}
fn into_each(self, mut accept: impl FnMut(i64) -> Result<()>) -> Result<()> {
match self {
Self::Empty => Ok(()),
Self::One(value) => accept(value),
Self::Many(values) => {
for value in values {
accept(value)?;
}
Ok(())
}
}
}
}
impl<'s> Spilled<'s> {
fn new(reader: Reader<'s>, types: Vec<LogicalType>) -> Self {
let columns = vec![Vec::new(); types.len()];
Self { reader, types, row: Vec::new(), columns }
}
fn next(&mut self, pass: &Aggregate<'_>) -> Result<Option<Rows>> {
let Self { reader, types, row, columns } = self;
for column in columns.iter_mut() {
column.clear();
}
let mut rows = 0;
while rows < VECTOR_SIZE && reader.next_into(row)? {
for (at, value) in row.iter_mut().enumerate() {
columns[at].push(std::mem::replace(value, Value::Null));
}
rows += 1;
}
if rows == 0 {
return Ok(None);
}
let mut built = Vec::with_capacity(columns.len());
for (values, ty) in columns.iter().zip(&*types) {
built.push(Vector::from_values(ty.clone(), values)?);
}
let mut taking = built.into_iter();
let keys: Vec<Vector> = taking.by_ref().take(pass.keys.len()).collect();
let mut arguments = Vec::with_capacity(pass.calls.len());
for call in &pass.calls {
arguments.push(taking.by_ref().take(call.args.len()).collect());
}
let mut filters = Vec::with_capacity(pass.calls.len());
for call in &pass.calls {
filters.push(if call.filter.is_some() { taking.next() } else { None });
}
Ok(Some(Rows { keys, arguments, filters, rows, marked: None }))
}
}
fn put_away(file: &mut Spill, seen: &Rows, row: usize, away: &mut Vec<Value>) -> Result<()> {
let columns = seen
.keys
.iter()
.chain(seen.arguments.iter().flatten())
.chain(seen.filters.iter().flatten());
away.truncate(seen.width());
for (at, column) in columns.enumerate() {
match away.get_mut(at) {
Some(slot) => set(slot, column, row)?,
None => away.push(column.try_value_at(row)?),
}
}
file.write(away)
}
const PASSES: u64 = 64;
fn hopeless(file: &Spill, groups: usize) -> Result<()> {
let left = file.rows() / width_of(groups).max(1);
if left > PASSES {
return Err(Error::out_of_memory(format!(
"the memory limit leaves room for {groups} groups at a time and {} rows have already \
gone to a spill file, which is more passes over it than this will finish in",
file.rows()
)));
}
Ok(())
}
fn crowded(memory: &Memory) -> bool {
match memory.limit() {
Some(limit) => memory.used() >= limit / 2,
None => false,
}
}
fn copies_overlap(folded: u64, groups: u64) -> bool {
groups > 0 && folded.saturating_mul(5) >= groups.saturating_mul(7)
}
fn keys_arrive_together(spreading: &Spreading) -> bool {
spreading.runs > 0 && spreading.split_rows >= spreading.runs.saturating_mul(2)
}
fn fill(key: &mut Key, columns: &[Vector], row: usize) -> Result<()> {
key.0.truncate(columns.len());
for (at, column) in columns.iter().enumerate() {
match key.0.get_mut(at) {
Some(slot) => set(slot, column, row)?,
None => key.0.push(column.try_value_at(row)?),
}
}
Ok(())
}
fn set(slot: &mut Value, column: &Vector, row: usize) -> Result<()> {
if let (Value::Varchar(buffer), Some(text)) = (&mut *slot, column.try_text_at(row)?) {
buffer.clear();
buffer.push_str(text);
return Ok(());
}
*slot = column.try_value_at(row)?;
Ok(())
}
impl Sink for Aggregate<'_> {
type Local = Partitioned;
fn finalize_degree(&self, ceiling: usize) -> usize {
if self.groups.is_empty() { 1 } else { ceiling }
}
fn weight(&self) -> usize {
let keys = self.keys.iter().map(|&key| {
let ty = self.plan.expr_type(key);
if ty.physical() == PhysicalType::Varlen || ty.is_nested() { 4 } else { 1 }
});
keys.sum::<usize>() + self.inputs.passes()
}
fn local(&self) -> Partitioned {
self.started.fetch_add(1, Ordering::Relaxed);
Partitioned {
mixed: group_mixed::Local::new(&self.memory),
counted: group_count::Local::new(&self.memory),
ranged: group_ranged::Local::new(&self.memory),
grouped_distinct: group_distinct::Local::new(&self.memory),
encoded: false,
encoded_records: EncodedRecords::new(self.keys.len()),
encoded_memory: self.memory.reservation(),
radix_distinct: false,
radix_distinct_records: (0..RADIX_PARTITIONS)
.map(|_| BigIntDistinctPartition::default())
.collect(),
radix_distinct_memory: self.memory.reservation(),
fixed: false,
fixed_records: (0..RADIX_PARTITIONS).map(|_| FixedRun::default()).collect(),
fixed_blocks: FixedBlocks::default(),
fixed_memory: self.memory.reservation(),
dense: false,
dense_codes: (0..DENSE_PARTITIONS).map(|_| Blocks::default()).collect(),
dense_nulls: 0,
dense_memory: self.memory.reservation(),
single: Some(self.start()),
closed: None,
ran: Vec::new(),
ran_memory: self.memory.reservation(),
installed: false,
expressions: self.inputs.scratch(),
spreading: Spreading::new(),
own: (0..RADIX_PARTITIONS).map(|_| None).collect(),
folded: 0,
}
}
fn parallel(&self) -> bool {
self.max_groups.is_none() || (!self.alone && !(self.count_only && self.keys.len() == 1))
}
fn sink(&self, chunk: &Chunk, local: &mut Partitioned) -> Result<Progress> {
let Partitioned {
mixed,
counted,
ranged,
grouped_distinct,
encoded,
encoded_records,
encoded_memory,
radix_distinct,
radix_distinct_records,
radix_distinct_memory,
fixed,
fixed_records,
fixed_memory,
fixed_blocks,
dense,
dense_codes,
dense_nulls,
dense_memory,
single,
closed,
ran,
ran_memory,
installed,
expressions,
spreading,
own,
folded,
} = local;
let rows = self.read(chunk, expressions)?;
if self.ranged_counts() {
return self.count_ranged(&rows, ranged);
}
if self.mixed_top_count() {
let [group] = rows.keys.as_slice() else {
return Err(Error::internal("a mixed radix exchange received the wrong key width"));
};
let [sum] = rows.arguments[0].as_slice() else {
return Err(Error::internal("a mixed radix exchange received no SUM argument"));
};
let [mean] = rows.arguments[2].as_slice() else {
return Err(Error::internal("a mixed radix exchange received no AVG argument"));
};
let [user] = rows.arguments[3].as_slice() else {
return Err(Error::internal(
"a mixed radix exchange received no distinct argument",
));
};
let buffered = group_mixed::Exchange::buffer(
&self.mixed,
&self.memory,
[group, sum, mean, user],
rows.rows,
mixed,
);
buffered?;
return Ok(Progress::More);
}
if self.counted_top_count() {
let [key] = rows.keys.as_slice() else {
return Err(Error::internal(
"a counted radix exchange received the wrong key width",
));
};
group_count::Exchange::buffer(
&self.counted,
self.plan.expr_type(self.keys[0]),
key,
rows.rows,
counted,
)?;
return Ok(Progress::More);
}
if self.grouped_distinct_top_count() {
if rows.keys.len() != self.keys.len() {
return Err(Error::internal(
"a grouped distinct exchange received the wrong key width",
));
}
let Some(user) = rows.arguments.first().and_then(|arguments| arguments.first()) else {
return Err(Error::internal("a grouped distinct exchange received no argument"));
};
let keys: Vec<group_distinct::Key<'_>> = rows
.keys
.iter()
.zip(&self.keys)
.map(|(vector, &key)| {
let kind = self.plan.expr_type(key);
let codes = if kind == &LogicalType::Varchar {
match vector.stable_dictionary_parts() {
Some((codes, dictionary))
if !dictionary.validity().has_nulls(dictionary.len()) =>
{
group_distinct::Codes::Dictionary(codes, dictionary)
}
_ => group_distinct::Codes::Loose,
}
} else {
group_distinct::Codes::Signed
};
group_distinct::Key { vector, kind, codes }
})
.collect();
let timing = stage::Timing::start(Stage::Scatter);
let buffered = group_distinct::Exchange::buffer(
&self.grouped_distinct,
&keys,
user,
rows.rows,
grouped_distinct,
);
timing.stop(0);
if buffered? {
return Ok(Progress::More);
}
}
if self.encoded_top_count() {
let timing = stage::Timing::start(Stage::Scatter);
let buffered = self.buffer_encoded_count(&rows, encoded_records, encoded_memory);
timing.stop(0);
if buffered? {
*encoded = true;
return Ok(Progress::More);
}
}
if self.radix_distinct_count {
let timing = stage::Timing::start(Stage::Scatter);
let buffered =
self.buffer_bigint_distinct(&rows, radix_distinct_records, radix_distinct_memory);
timing.stop(0);
buffered?;
*radix_distinct = true;
return Ok(Progress::More);
}
if self.fixed_top_count() {
let timing = stage::Timing::start(Stage::Scatter);
let buffered = self.buffer_fixed(&rows, fixed_records, fixed_memory, fixed_blocks);
timing.stop(0);
buffered?;
*fixed = true;
return Ok(Progress::More);
}
if self.count_only
&& self.keys.len() == 1
&& let [key] = rows.keys.as_slice()
&& let Some((codes, dictionary)) = key.stable_dictionary_parts()
{
let state = self.dense.get_or_init(|| DenseCount {
dictionary: Arc::clone(dictionary),
partitions: (0..DENSE_PARTITIONS)
.map(|_| Mutex::new(DensePartition::default()))
.collect(),
held: Mutex::new(Vec::new()),
});
if !Arc::ptr_eq(&state.dictionary, dictionary) {
return Err(Error::internal(
"one stable dictionary aggregate received two code spaces",
));
}
let validity = key.validity();
let before = dense_codes.iter().map(Blocks::footprint).sum::<usize>();
if !validity.has_nulls(rows.rows) && !dictionary.validity().has_nulls(dictionary.len())
{
for &code in &codes[..rows.rows] {
if code as usize >= dictionary.len() {
return Err(Error::internal("a stable dictionary code is out of range"));
}
dense_codes[code as usize % DENSE_PARTITIONS].push(code);
}
} else {
for (row, &code) in codes.iter().enumerate().take(rows.rows) {
if key.is_null_at(row) {
*dense_nulls += 1;
} else {
let code = code as usize;
if code >= dictionary.len() {
return Err(Error::internal(
"a stable dictionary code is out of range",
));
}
dense_codes[code % DENSE_PARTITIONS].push(code as u32);
}
}
}
let after = dense_codes.iter().map(Blocks::footprint).sum::<usize>();
dense_memory.grow(width_of(after.saturating_sub(before)))?;
*dense = true;
return Ok(Progress::More);
}
if self.closes()
&& let Some((from, to)) = interior(&rows.keys[0], rows.rows, self.grouped)
&& let (Ok(head), Ok(tail)) = (rows.slice(0, from), rows.slice(to, rows.rows - to))
{
if self.closes_by_run() {
let timing = stage::Timing::start(Stage::Fold);
let answered = self.close_runs(&rows, from, to);
timing.stop(0);
if let Some(answered) = answered? {
ran_memory.grow(width_of(answered.footprint()))?;
ran.push(answered);
self.open(&head, single, installed, spreading, own, folded)?;
self.open(&tail, single, installed, spreading, own, folded)?;
return Ok(Progress::More);
}
}
let building = closed.get_or_insert_with(|| self.shut());
let timing = stage::Timing::start(Stage::Fold);
let done = self.fold(&rows, building, None, Some((from, to)));
timing.stop(0);
done?;
self.open(&head, single, installed, spreading, own, folded)?;
self.open(&tail, single, installed, spreading, own, folded)?;
return Ok(Progress::More);
}
self.open(&rows, single, installed, spreading, own, folded)?;
Ok(Progress::More)
}
fn combine(&self, local: Partitioned) -> Result<()> {
let Partitioned {
mixed,
counted,
ranged,
grouped_distinct,
encoded,
mut encoded_records,
encoded_memory,
radix_distinct,
mut radix_distinct_records,
radix_distinct_memory,
fixed,
mut fixed_records,
fixed_memory,
dense,
mut dense_codes,
dense_nulls,
dense_memory,
single,
closed,
mut ran,
ran_memory,
mut spreading,
mut own,
folded,
..
} = local;
self.drain(folded, &mut spreading, &mut own)?;
if !ran.is_empty() {
let mut built = self.built.lock().map_err(poisoned)?;
built.chunks.append(&mut ran);
built.held.push(ran_memory);
}
if let Some(closed) = closed {
if let Some(error) = closed.failure {
return Err(error);
}
let mut chunks = Vec::new();
let mut held = self.memory.reservation();
if self.finish(closed, &mut chunks, &mut held)?.is_some() {
return Err(Error::internal("a closed group went to a spill file"));
}
let mut built = self.built.lock().map_err(poisoned)?;
built.chunks.append(&mut chunks);
built.held.push(held);
}
if mixed.used() {
let state = self.mixed.get().expect("a mixed exchange exists after its sink");
state.combine(mixed)?;
self.built.lock().map_err(poisoned)?.instances += 1;
return Ok(());
}
if ranged.used() {
let state = self.ranged.get().expect("a ranged exchange exists after its sink");
state.combine(ranged)?;
self.built.lock().map_err(poisoned)?.instances += 1;
return Ok(());
}
if counted.used() {
let state = self.counted.get().expect("a counted exchange exists after its sink");
state.combine(counted)?;
self.built.lock().map_err(poisoned)?.instances += 1;
return Ok(());
}
if grouped_distinct.used() {
let state = self
.grouped_distinct
.get()
.and_then(Option::as_ref)
.expect("a grouped distinct exchange exists after its sink buffered a chunk");
state.combine(grouped_distinct)?;
self.built.lock().map_err(poisoned)?.instances += 1;
return Ok(());
}
if encoded {
let state = self
.encoded_count
.get()
.and_then(Option::as_ref)
.expect("an encoded exchange exists after an encoded sink");
match (state, &mut encoded_records) {
(EncodedExchange::Two(state), EncodedRecords::Two(runs)) => {
state.take(runs, encoded_memory)?;
}
(
EncodedExchange::Three { tucked, wide },
EncodedRecords::Three { tucked: tucked_runs, wide: wide_runs },
) => {
tucked.take(tucked_runs, encoded_memory)?;
wide.take(wide_runs, self.memory.reservation())?;
}
_ => return Err(Error::internal("an encoded count exchange changed key width")),
}
self.built.lock().map_err(poisoned)?.instances += 1;
return Ok(());
}
if radix_distinct {
let state = self
.bigint_distinct
.get()
.expect("a distinct exchange exists after a distinct sink");
for (partition, rows) in radix_distinct_records.iter_mut().enumerate() {
if rows.rows.is_empty() {
continue;
}
let run = std::mem::take(&mut rows.rows);
state.partitions[partition].lock().map_err(poisoned)?.runs.push(run);
}
state.held.lock().map_err(poisoned)?.push(radix_distinct_memory);
self.built.lock().map_err(poisoned)?.instances += 1;
return Ok(());
}
if fixed {
let state = self.fixed.get().expect("fixed exchange exists after a fixed sink");
for (partition, rows) in fixed_records.iter_mut().enumerate() {
if rows.rows.is_empty() {
continue;
}
let run = std::mem::take(rows);
state.partitions[partition].lock().map_err(poisoned)?.runs.push(run);
}
state.held.lock().map_err(poisoned)?.push(fixed_memory);
self.built.lock().map_err(poisoned)?.instances += 1;
return Ok(());
}
if dense {
let state = self.dense.get().expect("dense state exists after a dense sink");
for (partition, run) in dense_codes.iter_mut().enumerate() {
if run.is_empty() && (partition != 0 || dense_nulls == 0) {
continue;
}
let mut shared = state.partitions[partition].lock().map_err(poisoned)?;
if partition == 0 {
shared.nulls += dense_nulls;
}
if !run.is_empty() {
shared.runs.push(std::mem::take(run));
}
}
state.held.lock().map_err(poisoned)?.push(dense_memory);
self.built.lock().map_err(poisoned)?.instances += 1;
return Ok(());
}
self.deposit(&mut own, &mut spreading)?;
let mut built = self.built.lock().map_err(poisoned)?;
built.instances += 1;
let Some(arriving) = single else { return Ok(()) };
if let Some(error) = arriving.failure {
return Err(error);
}
if built.partitioning {
drop(built);
self.hand(arriving, &mut spreading, &mut own)?;
return self.deposit(&mut own, &mut spreading);
}
drop(built);
let mut arriving = arriving;
loop {
let mut built = self.built.lock().map_err(poisoned)?;
if built.partitioning {
drop(built);
self.hand(arriving, &mut spreading, &mut own)?;
return self.deposit(&mut own, &mut spreading);
}
let mut kept = self.merged[0].lock().map_err(poisoned)?;
let Some(waiting) = kept.table.take() else {
kept.table = Some(arriving);
return Ok(());
};
if arriving.over.is_some() || waiting.over.is_some() {
built.partitioning = true;
drop(kept);
drop(built);
self.hand_over(waiting, &mut spreading)?;
return self.hand_over(arriving, &mut spreading);
}
drop(kept);
drop(built);
let (from, mut into) = if waiting.groups > arriving.groups {
(arriving, waiting)
} else {
(waiting, arriving)
};
self.merge(from, &mut into)?;
arriving = into;
}
}
fn finalize(&self, threads: &Lease<'_>) -> Result<()> {
if let Some(mixed) = self.mixed.get() {
let chunks = mixed.finish(
threads,
self.top_counts.expect("a mixed exchange has a TopN bound").0,
&self.memory,
)?;
return self.out.fill(chunks);
}
if let Some(ranged) = self.ranged.get() {
let chunks = ranged.finish(&self.memory)?;
return self.out.fill(chunks);
}
if let Some(counted) = self.counted.get() {
let bound = self.top_counts.expect("a counted exchange has a TopN bound").0;
let degree = fixed_degree(counted.records()?, threads);
let chunks = counted.finish(threads, degree, bound, &self.memory)?;
return self.out.fill(chunks);
}
if let Some(Some(distinct)) = self.grouped_distinct.get() {
let chunks = distinct.finish(
threads,
self.top_counts.expect("a grouped distinct exchange has a TopN bound").0,
&self.memory,
)?;
return self.out.fill(chunks);
}
if let Some(Some(encoded)) = self.encoded_count.get() {
return match encoded {
EncodedExchange::Two(encoded) => {
let chunks = self.finalize_encoded(encoded, threads)?;
self.out.fill(chunks)
}
EncodedExchange::Three { tucked, wide } => {
let mut chunks = self.finalize_encoded(tucked, threads)?;
chunks.append(&mut self.finalize_encoded(wide, threads)?);
self.out.fill(chunks)
}
};
}
if let Some(distinct) = self.bigint_distinct.get() {
let next = AtomicUsize::new(0);
let slots: Vec<Mutex<Option<Result<i64>>>> =
(0..RADIX_PARTITIONS).map(|_| Mutex::new(None)).collect();
let input = distinct
.partitions
.iter()
.map(|partition| {
partition
.lock()
.map(|held| held.runs.iter().map(Vec::len).sum::<usize>())
.map_err(poisoned)
})
.sum::<Result<usize>>()?;
let degree = degree_for(input, threads);
together(threads, degree, &|| {
finish_bigint_distinct(&next, &slots, distinct, &self.memory);
})?;
let mut total = 0_i64;
for (at, slot) in slots.iter().enumerate() {
let count = slot.lock().map_err(poisoned)?.take().unwrap_or_else(|| {
Err(Error::internal(format!("nothing finished distinct radix partition {at}")))
})?;
total = total
.checked_add(count)
.ok_or_else(|| Error::out_of_range("COUNT(DISTINCT BIGINT) overflowed"))?;
}
let mut held = distinct.held.lock().map_err(poisoned)?;
held.clear();
let values = [vec![Value::BigInt(total)]];
let mut output = self.memory.reservation();
let chunks = rows::chunks(&[LogicalType::BigInt], &values, &mut output)?;
held.push(output);
drop(held);
return self.out.fill(chunks);
}
if let Some(fixed) = self.fixed.get() {
let bound = self.top_counts.expect("a fixed exchange has a TopN bound").0;
let next = AtomicUsize::new(0);
let slots: Vec<Mutex<Option<Result<Part>>>> =
(0..RADIX_PARTITIONS).map(|_| Mutex::new(None)).collect();
let input = fixed
.partitions
.iter()
.map(|partition| {
partition
.lock()
.map(|runs| runs.runs.iter().map(|run| run.rows.len()).sum::<usize>())
.map_err(poisoned)
})
.sum::<Result<usize>>()?;
let degree = fixed_degree(input, threads);
together(threads, degree, &|| {
finish_fixed(&next, &slots, fixed, bound, &self.calls, &self.memory);
})?;
let mut parts = Vec::with_capacity(slots.len());
for (at, slot) in slots.iter().enumerate() {
parts.push(slot.lock().map_err(poisoned)?.take().unwrap_or_else(|| {
Err(Error::internal(format!("nothing finished fixed radix partition {at}")))
})?);
}
let mut chunks = Vec::new();
let mut held = fixed.held.lock().map_err(poisoned)?;
held.clear();
for Part { chunks: mut part, held: charge } in parts {
chunks.append(&mut part);
held.push(charge);
}
drop(held);
return self.out.fill(chunks);
}
if let Some(dense) = self.dense.get() {
let mut working = self.memory.reservation();
working.grow(width_of(dense.dictionary.len() * size_of::<i64>()))?;
let types = self.schema.types();
let group_types = &types[..self.groups.len()];
let next = AtomicUsize::new(0);
let slots: Vec<Mutex<Option<Result<Vec<Chunk>>>>> =
(0..dense.partitions.len()).map(|_| Mutex::new(None)).collect();
let degree = threads.degree().clamp(1, slots.len().max(1));
together(threads, degree, &|| {
loop {
let at = next.fetch_add(1, Ordering::Relaxed);
let Some(partition) = dense.partitions.get(at) else { return };
let done = partition.lock().map_err(poisoned).and_then(|mut held| {
dense_partition(
&dense.dictionary,
at,
&mut held,
self.top_counts.map(|(bound, _)| bound),
&self.constants,
group_types,
)
});
if let Ok(mut slot) = slots[at].lock() {
*slot = Some(done);
}
}
})?;
let mut chunks = Vec::new();
for (at, slot) in slots.iter().enumerate() {
chunks.extend(slot.lock().map_err(poisoned)?.take().unwrap_or_else(|| {
Err(Error::internal(format!("nothing finished dense partition {at}")))
})?);
}
let mut output = self.memory.reservation();
let shared = dense.dictionary.footprint();
let output_bytes = chunks
.iter()
.map(Chunk::footprint)
.sum::<usize>()
.saturating_sub(shared.saturating_mul(chunks.len().saturating_sub(1)));
output.grow(width_of(output_bytes))?;
let mut held = dense.held.lock().map_err(poisoned)?;
held.clear();
working.release();
held.push(output);
drop(held);
return self.out.fill(chunks);
}
let mut built = self.built.lock().map_err(poisoned)?;
let degree = built.instances.min(threads.degree()).clamp(1, self.merged.len());
let closed =
if degree > 1 { self.close_together(threads, degree)? } else { self.close_in_turn()? };
for part in closed {
let Part { mut chunks, held } = part?;
built.chunks.append(&mut chunks);
built.held.push(held);
}
let chunks = std::mem::take(&mut built.chunks);
drop(built);
self.out.fill(chunks)
}
}
fn finish_encoded_count<S: Second>(
next: &AtomicUsize,
slots: &[Mutex<Option<Result<Part>>>],
encoded: &EncodedCountExchange<S>,
bound: usize,
memory: &Memory,
) {
loop {
let at = next.fetch_add(1, Ordering::Relaxed);
let Some(partition) = encoded.partitions.get(at) else {
return;
};
let done = partition.lock().map_err(poisoned).and_then(|mut rows| {
encoded_count_partition(
&mut rows,
&encoded.dictionary,
&encoded.leading,
encoded.code_bits,
bound,
memory,
)
});
if let Ok(mut slot) = slots[at].lock() {
*slot = Some(done);
}
}
}
#[inline]
fn encoded_slot<S: Second>(
buckets: &[u32],
groups: &EncodedCountPartition<S>,
row: EncodedCountRecord<S>,
valid: u8,
) -> std::result::Result<usize, usize> {
let mask = buckets.len() - 1;
let all_valid = groups.validity.is_empty();
let tag = slot_tag(u64::from(row.hash));
let mut at = row.hash as usize & mask;
loop {
let bucket = buckets[at];
let slot = bucket & SLOT_MASK;
if slot == EMPTY_SLOT {
return Err(at);
}
if bucket == tag | slot {
let slot = slot as usize;
let held = groups.rows[slot];
let held_valid = if all_valid { EncodedValid::ALL } else { groups.validity[slot] };
if held.hash == row.hash
&& held.first == row.first
&& held.second == row.second
&& held.third == row.third
&& held_valid == valid
{
return Ok(slot);
}
}
at = (at + 1) & mask;
}
}
#[inline]
fn encoded_split(hash: u32, splits: usize) -> usize {
(hash >> 16) as usize & (splits - 1)
}
const ENCODED_SPLITS: usize = 256;
fn encoded_count_partition<S: Second>(
runs: &mut EncodedCountRuns<S>,
dictionary: &Vector,
leading: &[LogicalType],
code_bits: u32,
bound: usize,
memory: &Memory,
) -> Result<Part> {
let keys = leading.len() + 1;
if !(2..=3).contains(&keys) {
return Err(Error::internal("an encoded count partition has an unsupported key width"));
}
let reserving = stage::Timing::start(Stage::Reserve);
let total: usize = runs.runs.iter().map(EncodedCountRun::<S>::len).sum();
let splits = (total / FIXED_SPLIT_ROWS).max(1).next_power_of_two().min(ENCODED_SPLITS);
let share = total.div_ceil(splits);
let share = (share + share.isqrt() * 4).min(total);
let capacity = share.saturating_mul(2).max(64).next_power_of_two();
let mut working = memory.reservation();
working.grow(width_of(
capacity * size_of::<u32>()
+ share * size_of::<i64>()
+ total * size_of::<EncodedCountRecord<S>>(),
))?;
let mut parts: Vec<EncodedCountPartition<S>> = (0..splits)
.map(|_| EncodedCountPartition {
rows: Vec::with_capacity(share),
validity: Vec::new(),
weights: Vec::new(),
})
.collect();
reserving.stop(0);
let timing = stage::Timing::start(Stage::Merge);
for run in std::mem::take(&mut runs.runs) {
let all_valid = run.group_validity.is_empty();
let mut source = 0;
for (block, weights) in run.groups.slices().zip(run.weights.slices()) {
for (&row, &weight) in block.iter().zip(weights) {
let valid = if all_valid { EncodedValid::ALL } else { run.group_validity[source] };
parts[encoded_split(row.hash, splits)].push_weighted(row, valid, weight);
source += 1;
}
}
let all_valid = run.pending_validity.is_empty();
let mut source = 0;
for block in run.pending.slices() {
for &row in block {
let valid =
if all_valid { EncodedValid::ALL } else { run.pending_validity[source] };
let weight = run.pending_weight(source);
parts[encoded_split(row.hash, splits)].push_weighted(row, valid, weight);
source += 1;
}
}
}
timing.stop(0);
let timing = stage::Timing::start(Stage::Fold);
let mut buckets: Vec<u32> = Vec::with_capacity(capacity);
let mut counts: Vec<i64> = Vec::with_capacity(share);
let mut output: Vec<(i64, Vec<Value>)> = Vec::new();
for mut partition in parts {
let rows = partition.rows.len();
buckets.clear();
buckets.resize(rows.saturating_mul(2).max(64).next_power_of_two(), EMPTY_SLOT);
counts.clear();
let all_valid = partition.validity.is_empty();
for source in 0..rows {
let row = partition.rows[source];
let valid = if all_valid { EncodedValid::ALL } else { partition.validity[source] };
let weight = partition.weight(source);
let slot = match encoded_slot(&buckets, &partition, row, valid) {
Ok(slot) => slot,
Err(bucket) => {
let slot = counts.len();
buckets[bucket] = bucket_for(
slot,
u64::from(row.hash),
"an encoded radix partition is too large",
)?;
partition.rows[slot] = row;
if !all_valid {
partition.validity[slot] = valid;
}
counts.push(0);
slot
}
};
counts[slot] = counts[slot]
.checked_add(weight)
.ok_or_else(|| Error::out_of_range("a grouped COUNT overflowed BIGINT"))?;
}
let best = largest(counts.len(), bound, |slot| counts[slot]);
for slot in best {
let key = partition.rows[slot];
let (second, code) = key.second.keys(key.third, code_bits);
let valid = if all_valid { EncodedValid::ALL } else { partition.validity[slot] };
let first = if valid & EncodedValid::FIRST != 0 {
signed_value(&leading[0], key.first)?
} else {
Value::Null
};
let third = if valid & EncodedValid::THIRD != 0 {
dictionary.try_value_at(code as usize)?
} else {
Value::Null
};
let mut row = Vec::with_capacity(keys + 1);
row.push(first);
if keys == 3 {
row.push(if valid & EncodedValid::SECOND != 0 {
signed_value(&leading[1], second)?
} else {
Value::Null
});
}
row.push(third);
row.push(Value::BigInt(counts[slot]));
output.push((counts[slot], row));
}
}
timing.stop(0);
let timing = stage::Timing::start(Stage::Emit);
output.sort_by_key(|(count, _)| std::cmp::Reverse(*count));
output.truncate(bound);
let output = output.into_iter().map(|(_, row)| row).collect::<Vec<_>>();
let mut held = memory.reservation();
let mut types = leading.to_vec();
types.push(LogicalType::Varchar);
types.push(LogicalType::BigInt);
let chunks = rows::chunks(&types, &output, &mut held)?;
timing.stop(0);
Ok(Part { chunks, held })
}
fn finish_bigint_distinct(
next: &AtomicUsize,
slots: &[Mutex<Option<Result<i64>>>],
distinct: &BigIntDistinctExchange,
memory: &Memory,
) {
loop {
let at = next.fetch_add(1, Ordering::Relaxed);
let Some(partition) = distinct.partitions.get(at) else {
return;
};
let done = partition
.lock()
.map_err(poisoned)
.and_then(|mut rows| bigint_distinct_partition(&mut rows, memory));
if let Ok(mut slot) = slots[at].lock() {
*slot = Some(done);
}
}
}
#[inline]
fn scatter_bigint(partitions: &mut [BigIntDistinctPartition], shift: u32, value: i64) {
let hash = spread(mix(0, value as u64));
partitions[(hash >> shift) as usize].rows.push(value);
}
fn distinct_table_bytes(capacity: usize) -> usize {
capacity * size_of::<i64>() + capacity.div_ceil(8)
}
fn regrow_distinct(slots: &mut Vec<i64>, filled: &mut Vec<u64>) {
let capacity = slots.len() * 2;
let mask = capacity - 1;
let mut next = vec![0_i64; capacity];
let mut taken = vec![0_u64; capacity.div_ceil(64)];
for (from, &value) in slots.iter().enumerate() {
if filled[from / 64] & (1_u64 << (from % 64)) == 0 {
continue;
}
let mut at = spread(mix(0, value as u64)) as usize & mask;
while taken[at / 64] & (1_u64 << (at % 64)) != 0 {
at = (at + 1) & mask;
}
taken[at / 64] |= 1_u64 << (at % 64);
next[at] = value;
}
*slots = next;
*filled = taken;
}
fn bigint_distinct_partition(partition: &mut BigIntDistinctRuns, memory: &Memory) -> Result<i64> {
let held: usize = partition.runs.iter().map(Vec::len).sum();
let ceiling = held.saturating_mul(2).max(64).next_power_of_two();
let mut capacity = ceiling.min(1024);
let mut working = memory.reservation();
working.grow(width_of(distinct_table_bytes(capacity)))?;
let mut slots = vec![0_i64; capacity];
let mut filled = vec![0_u64; capacity.div_ceil(64)];
let mut mask = capacity - 1;
let mut limit = capacity / 2;
let mut unique = 0_usize;
let timing = stage::Timing::start(Stage::Fold);
for run in &partition.runs {
for &value in run {
let mut at = spread(mix(0, value as u64)) as usize & mask;
loop {
let bit = 1_u64 << (at % 64);
if filled[at / 64] & bit == 0 {
filled[at / 64] |= bit;
slots[at] = value;
unique += 1;
break;
}
if slots[at] == value {
break;
}
at = (at + 1) & mask;
}
if unique >= limit && capacity < ceiling {
working.grow(width_of(distinct_table_bytes(capacity)))?;
regrow_distinct(&mut slots, &mut filled);
capacity *= 2;
mask = capacity - 1;
limit = capacity / 2;
}
}
}
timing.stop(0);
i64::try_from(unique).map_err(|_| Error::out_of_range("COUNT(DISTINCT BIGINT) overflowed"))
}
fn finish_fixed(
next: &AtomicUsize,
slots: &[Mutex<Option<Result<Part>>>],
fixed: &FixedExchange,
bound: usize,
calls: &[Call],
memory: &Memory,
) {
loop {
let at = next.fetch_add(1, Ordering::Relaxed);
let Some(partition) = fixed.partitions.get(at) else {
return;
};
let done = partition
.lock()
.map_err(poisoned)
.and_then(|mut rows| fixed_partition(&mut rows, &fixed.keys, bound, calls, memory));
if let Ok(mut slot) = slots[at].lock() {
*slot = Some(done);
}
}
}
#[inline]
fn fixed_slot(
buckets: &[u32],
groups: &FixedPartition,
row: FixedRecord,
valid: u8,
) -> std::result::Result<usize, usize> {
const KEYS: u8 = FixedRecord::FIRST | FixedRecord::SECOND;
let mask = buckets.len() - 1;
let all_valid = groups.validity.is_empty();
let hash = fixed_hash(row, valid);
let tag = slot_tag(hash);
let mut at = hash as usize & mask;
loop {
let bucket = buckets[at];
let slot = bucket & SLOT_MASK;
if slot == EMPTY_SLOT {
return Err(at);
}
if bucket == tag | slot {
let slot = slot as usize;
let held = groups.rows[slot];
let held_valid = if all_valid { FixedRecord::ALL } else { groups.validity[slot] };
if held.first == row.first
&& held.second == row.second
&& held_valid & KEYS == valid & KEYS
{
return Ok(slot);
}
}
at = (at + 1) & mask;
}
}
const FIXED_SPLIT_ROWS: usize = 16_384;
#[inline]
fn fixed_split(hash: u64, splits: usize) -> usize {
(hash >> 40) as usize & (splits - 1)
}
fn fixed_partition(
runs: &mut FixedRuns,
keys: &[LogicalType; 2],
bound: usize,
calls: &[Call],
memory: &Memory,
) -> Result<Part> {
let reserving = stage::Timing::start(Stage::Reserve);
let total: usize = runs.runs.iter().map(|run| run.rows.len()).sum();
let splits = (total / FIXED_SPLIT_ROWS).max(1).next_power_of_two();
let share = total.div_ceil(splits);
let share = (share + share.isqrt() * 4).min(total);
let capacity = share.saturating_mul(2).max(64).next_power_of_two();
let mut working = memory.reservation();
working.grow(width_of(
capacity * size_of::<u32>()
+ share * size_of::<CompactNumeric>()
+ total * size_of::<FixedRecord>(),
))?;
let mut parts: Vec<FixedPartition> = (0..splits)
.map(|_| FixedPartition { rows: Vec::with_capacity(share), validity: Vec::new() })
.collect();
reserving.stop(0);
let timing = stage::Timing::start(Stage::Merge);
for run in std::mem::take(&mut runs.runs) {
let all_valid = run.validity.is_empty();
let mut source = 0;
for block in run.rows.slices() {
for &row in block {
let valid = if all_valid { FixedRecord::ALL } else { run.validity[source] };
parts[fixed_split(fixed_hash(row, valid), splits)].push(row, valid);
source += 1;
}
}
}
timing.stop(0);
let timing = stage::Timing::start(Stage::Fold);
let mut buckets: Vec<u32> = Vec::with_capacity(capacity);
let mut states: Vec<CompactNumeric> = Vec::with_capacity(share);
let mut output: Vec<(i64, Vec<Value>)> = Vec::new();
for mut partition in parts {
let rows = partition.rows.len();
let capacity = rows.saturating_mul(2).max(64).next_power_of_two();
buckets.clear();
buckets.resize(capacity, EMPTY_SLOT);
states.clear();
let mut overflow = HashMap::new();
let all_valid = partition.validity.is_empty();
for source in 0..rows {
let row = partition.rows[source];
let valid = if all_valid { FixedRecord::ALL } else { partition.validity[source] };
let slot = match fixed_slot(&buckets, &partition, row, valid) {
Ok(slot) => slot,
Err(bucket) => {
let slot = states.len();
buckets[bucket] = bucket_for(
slot,
fixed_hash(row, valid),
"a fixed radix partition is too large",
)?;
partition.rows[slot] = row;
if !all_valid {
partition.validity[slot] = valid;
}
states.push(CompactNumeric::default());
slot
}
};
states[slot].add(
slot,
(valid & FixedRecord::SUM != 0).then_some(row.sum),
(valid & FixedRecord::MEAN != 0).then_some(row.mean),
&mut overflow,
)?;
}
let best = largest(states.len(), bound, |slot| states[slot].count());
for slot in best {
let key = partition.rows[slot];
let valid = if all_valid { FixedRecord::ALL } else { partition.validity[slot] };
let state = &states[slot];
let (sum, mean) = state.totals(slot, &overflow);
output.push((
state.count(),
vec![
if valid & FixedRecord::FIRST != 0 {
signed_value(&keys[0], key.first)?
} else {
Value::Null
},
if valid & FixedRecord::SECOND != 0 {
signed_value(&keys[1], i64::from(key.second))?
} else {
Value::Null
},
Value::BigInt(state.count()),
Accumulator::exact_sum(sum, state.sum_seen(), &calls[1].returns).finish()?,
Accumulator::exact_avg(mean, state.mean_count, &calls[2].returns).finish()?,
],
));
}
}
timing.stop(0);
let timing = stage::Timing::start(Stage::Emit);
output.sort_by_key(|(count, _)| std::cmp::Reverse(*count));
output.truncate(bound);
let output = output.into_iter().map(|(_, row)| row).collect::<Vec<_>>();
let mut held = memory.reservation();
let types = [
keys[0].clone(),
keys[1].clone(),
LogicalType::BigInt,
calls[1].returns.clone(),
calls[2].returns.clone(),
];
let chunks = rows::chunks(&types, &output, &mut held)?;
timing.stop(0);
Ok(Part { chunks, held })
}
fn dense_partition(
dictionary: &Arc<Vector>,
number: usize,
partition: &mut DensePartition,
bound: Option<usize>,
constants: &[Option<Value>],
group_types: &[LogicalType],
) -> Result<Vec<Chunk>> {
let width = dictionary.len().saturating_add(DENSE_PARTITIONS - 1 - number) / DENSE_PARTITIONS;
let rows: usize = partition.runs.iter().map(Blocks::len).sum();
let mut groups: Vec<(u32, i64)> = if rows.saturating_mul(SPARSE_DENSE) < width {
let mut sorted: Vec<u32> =
partition.runs.iter().flat_map(Blocks::slices).flatten().copied().collect();
sorted.sort_unstable();
sorted.chunk_by(|left, right| left == right).map(|run| (run[0], run.len() as i64)).collect()
} else if u32::try_from(rows).is_ok() {
dense_counts::<u32>(&partition.runs, width, number, bound)
} else {
dense_counts::<i64>(&partition.runs, width, number, bound)
};
if let Some(bound) = bound.filter(|&bound| bound < groups.len()) {
let mut best = largest(groups.len(), bound, |slot| groups[slot].1);
best.sort_unstable();
groups = best.into_iter().map(|slot| groups[slot]).collect();
}
let mut chunks = Vec::new();
let mut codes = Vec::with_capacity(VECTOR_SIZE);
let mut counts = Vec::with_capacity(VECTOR_SIZE);
let mut valid = Vec::with_capacity(VECTOR_SIZE);
for (code, count) in groups {
codes.push(code);
counts.push(count);
valid.push(true);
if codes.len() == VECTOR_SIZE {
chunks.push(dense_chunk(dictionary, &codes, &counts, &valid, constants, group_types)?);
codes.clear();
counts.clear();
valid.clear();
}
}
if number == 0 && partition.nulls != 0 {
codes.push(0);
counts.push(partition.nulls);
valid.push(false);
}
if !codes.is_empty() {
chunks.push(dense_chunk(dictionary, &codes, &counts, &valid, constants, group_types)?);
}
Ok(chunks)
}
fn dense_counts<C>(
runs: &[Blocks<u32>],
width: usize,
number: usize,
bound: Option<usize>,
) -> Vec<(u32, i64)>
where
C: Copy + Default + PartialEq + std::ops::AddAssign + From<u8> + Into<i64>,
{
let mut dense = vec![C::default(); width];
for run in runs {
for block in run.slices() {
for &code in block {
dense[code as usize / DENSE_PARTITIONS] += C::from(1);
}
}
}
let code = |slot: usize| (slot * DENSE_PARTITIONS + number) as u32;
if let Some(bound) = bound {
let mut best = largest(width, bound, |slot| dense[slot].into());
best.retain(|&slot| dense[slot] != C::default());
best.sort_unstable();
return best.into_iter().map(|slot| (code(slot), dense[slot].into())).collect();
}
dense
.iter()
.enumerate()
.filter(|(_, count)| **count != C::default())
.map(|(slot, &count)| (code(slot), count.into()))
.collect()
}
fn dense_chunk(
dictionary: &Arc<Vector>,
codes: &[u32],
counts: &[i64],
valid: &[bool],
constants: &[Option<Value>],
group_types: &[LogicalType],
) -> Result<Chunk> {
let validity = Validity::from_iter(valid.len(), |row| valid[row]);
let key =
Vector::stable_dictionary(codes.to_vec(), Arc::clone(dictionary))?.with_validity(validity);
let mut key = Some(key);
let mut columns = Vec::with_capacity(constants.len() + 1);
for (constant, ty) in constants.iter().zip(group_types) {
columns.push(match constant {
Some(value) => Vector::constant(ty.clone(), value.clone(), codes.len()),
None => key
.take()
.ok_or_else(|| Error::internal("a dense count has more than one varying key"))?,
});
}
let counts = Vector::flat(LogicalType::BigInt, Data::Int64(counts.to_vec().into()))?;
columns.push(counts);
Chunk::with_rows(columns, codes.len())
}
#[derive(Debug)]
struct Part {
chunks: Vec<Chunk>,
held: Reservation,
}
fn degree_for(input: usize, threads: &Lease<'_>) -> usize {
pairs::finish_degree(input, RADIX_PARTITIONS.min(threads.degree()))
}
const FIXED_ROWS_PER_THREAD: usize = 8_192;
const FIXED_RAMP: usize = 16;
fn fixed_degree(input: usize, threads: &Lease<'_>) -> usize {
let quickly = input.div_ceil(FIXED_ROWS_PER_THREAD).min(FIXED_RAMP);
let slowly = input.div_ceil(pairs::ROWS_PER_EXTRA_THREAD);
quickly.max(slowly).clamp(1, RADIX_PARTITIONS.min(threads.degree()))
}
impl Aggregate<'_> {
fn finalize_encoded<S: Second>(
&self,
encoded: &EncodedCountExchange<S>,
threads: &Lease<'_>,
) -> Result<Vec<Chunk>> {
let next = AtomicUsize::new(0);
let slots: Vec<Mutex<Option<Result<Part>>>> =
(0..RADIX_PARTITIONS).map(|_| Mutex::new(None)).collect();
let input = encoded
.partitions
.iter()
.map(|partition| {
partition
.lock()
.map(|runs| runs.runs.iter().map(EncodedCountRun::<S>::len).sum::<usize>())
.map_err(poisoned)
})
.sum::<Result<usize>>()?;
let degree = degree_for(input, threads);
let bound = self.top_counts.expect("an encoded exchange has a TopN bound").0;
together(threads, degree, &|| {
finish_encoded_count(&next, &slots, encoded, bound, &self.memory);
})?;
let mut parts = Vec::with_capacity(slots.len());
for (at, slot) in slots.iter().enumerate() {
parts.push(slot.lock().map_err(poisoned)?.take().unwrap_or_else(|| {
Err(Error::internal(format!("nothing finished encoded radix partition {at}")))
})?);
}
let mut chunks = Vec::new();
let mut held = encoded.held.lock().map_err(poisoned)?;
held.clear();
for Part { chunks: mut part, held: charge } in parts {
chunks.append(&mut part);
held.push(charge);
}
drop(held);
Ok(chunks)
}
fn close_in_turn(&self) -> Result<Vec<Result<Part>>> {
Ok((0..self.merged.len()).map(|at| self.close(at)).collect())
}
fn close_together(&self, threads: &Lease<'_>, degree: usize) -> Result<Vec<Result<Part>>> {
let next = AtomicUsize::new(0);
let slots: Vec<Mutex<Option<Result<Part>>>> =
(0..self.merged.len()).map(|_| Mutex::new(None)).collect();
together(threads, degree, &|| self.closing(&next, &slots))?;
let mut closed = Vec::with_capacity(slots.len());
for (at, slot) in slots.into_iter().enumerate() {
closed.push(slot.into_inner().map_err(poisoned)?.unwrap_or_else(|| {
Err(Error::internal(format!("nothing finished partition {at} of an aggregate")))
}));
}
Ok(closed)
}
fn closing(&self, next: &AtomicUsize, slots: &[Mutex<Option<Result<Part>>>]) {
loop {
let at = next.fetch_add(1, Ordering::Relaxed);
let Some(slot) = slots.get(at) else { return };
let done = self.close(at);
if let Ok(mut slot) = slot.lock() {
*slot = Some(done);
}
}
}
fn close(&self, at: usize) -> Result<Part> {
let mut part = Part { chunks: Vec::new(), held: self.memory.reservation() };
let mut partition = self.merged[at].lock().map_err(poisoned)?;
let Partition { table, carried, pending } = &mut *partition;
let mut kept = table.take();
for arriving in std::mem::take(pending) {
match &mut kept {
Some(into) => self.merge(arriving, into)?,
None => kept = Some(arriving),
}
}
let Some(kept) = kept else {
debug_assert!(carried.is_none(), "nothing is set aside from a partition with no table");
return Ok(part);
};
let Part { chunks, held } = &mut part;
let mut left = self.finish(kept, chunks, held)?;
if left.is_none()
&& let Some(whole) = carried.take()
{
left = self.finish(whole, chunks, held)?;
}
while let Some(mut file) = left {
left = self.again(&mut file, carried.take(), chunks, held)?;
}
Ok(part)
}
}
fn nulls_are_in_the_mask(column: &Vector) -> bool {
!matches!(column.form(), Form::Dictionary | Form::Rle)
}
#[derive(Debug, Clone, Copy)]
enum Signed<'a> {
Int8(&'a [i8]),
Int16(&'a [i16]),
Int32(&'a [i32]),
Int64(&'a [i64]),
}
impl<'a> Signed<'a> {
fn of(column: &'a Vector, rows: usize) -> Option<Self> {
if column.validity().has_nulls(rows) {
return None;
}
match column.data()? {
Data::Int8(values) => values.get(..rows).map(Signed::Int8),
Data::Int16(values) => values.get(..rows).map(Signed::Int16),
Data::Int32(values) => values.get(..rows).map(Signed::Int32),
Data::Int64(values) => values.get(..rows).map(Signed::Int64),
_ => None,
}
}
#[inline]
fn at(self, row: usize) -> i64 {
match self {
Self::Int8(values) => i64::from(values[row]),
Self::Int16(values) => i64::from(values[row]),
Self::Int32(values) => i64::from(values[row]),
Self::Int64(values) => values[row],
}
}
}
fn signed_key(ty: &LogicalType) -> bool {
matches!(
ty,
LogicalType::TinyInt | LogicalType::SmallInt | LogicalType::Integer | LogicalType::BigInt
)
}
fn narrow_key(ty: &LogicalType) -> bool {
matches!(ty, LogicalType::TinyInt | LogicalType::SmallInt | LogicalType::Integer)
}
fn fixed_width(ty: &LogicalType) -> bool {
matches!(
ty,
LogicalType::Boolean
| LogicalType::TinyInt
| LogicalType::SmallInt
| LogicalType::Integer
| LogicalType::BigInt
| LogicalType::HugeInt
| LogicalType::UTinyInt
| LogicalType::USmallInt
| LogicalType::UInteger
| LogicalType::UBigInt
| LogicalType::UHugeInt
| LogicalType::Float
| LogicalType::Double
| LogicalType::Decimal { .. }
| LogicalType::Uuid
| LogicalType::Date
| LogicalType::Time
| LogicalType::Timestamp
| LogicalType::TimestampS
| LogicalType::TimestampMs
| LogicalType::TimestampNs
| LogicalType::TimestampTz
)
}
pub(crate) fn signed_value(ty: &LogicalType, value: i64) -> Result<Value> {
match ty {
LogicalType::TinyInt => i8::try_from(value).map(Value::TinyInt).map_err(|_| too_wide(ty)),
LogicalType::SmallInt => {
i16::try_from(value).map(Value::SmallInt).map_err(|_| too_wide(ty))
}
LogicalType::Integer => i32::try_from(value).map(Value::Integer).map_err(|_| too_wide(ty)),
LogicalType::BigInt => Ok(Value::BigInt(value)),
_ => Err(Error::internal(format!("{ty} is not a signed integer group key"))),
}
}
fn too_wide(ty: &LogicalType) -> Error {
Error::internal(format!("a radix group key does not fit back into {ty}"))
}
fn charge(into: &mut Building, grown: u64) -> Result<()> {
into.containers.grow(grown)?;
rows::capacity(into.table.owned(), &mut into.charged_keys, &mut into.scratch)?;
let now =
tables(&into.table, &into.states, &into.counts, &into.compact, &into.overflow, &into.seen);
rows::capacity(now, &mut into.charged, &mut into.containers)
}
struct Folding<'a> {
count_only: bool,
calls: usize,
distinct: &'a [bool],
taken: &'a [Accumulator],
tallies: &'a [i64],
compact: &'a [CompactNumeric],
overflow: &'a HashMap<usize, (i128, i128)>,
watched: &'a mut [DistinctSet],
}
fn merge_slot(
from: &mut Folding<'_>,
slot: usize,
target: usize,
into: &mut Building,
) -> Result<u64> {
if let Some(arriving) = from.compact.get(slot) {
let kept = &mut into.compact[target];
kept.combine(target, arriving, slot, from.overflow, &mut into.overflow)?;
return Ok(0);
}
if from.count_only {
into.counts[target] += from.tallies[slot];
return Ok(0);
}
let calls = from.calls;
let mut aside = 0;
for at in 0..calls {
if !from.distinct[at] {
into.states[target * calls + at].combine(&from.taken[slot * calls + at])?;
continue;
}
let arriving = std::mem::replace(
&mut from.watched[slot * calls + at],
DistinctSet::Row(RowSet::default()),
);
let state = &mut into.states[target * calls + at];
match (&mut into.seen[target * calls + at], arriving) {
(DistinctSet::BigInt(kept), DistinctSet::BigInt(arriving)) => {
arriving.into_each(|value| {
if kept.insert(value) {
aside += width_of(size_of::<i64>() * 2);
state.update(&[Value::BigInt(value)])?;
}
Ok(())
})?;
}
(DistinctSet::Row(kept), DistinctSet::Row(arriving)) => {
for key in arriving {
if kept.contains(&key) {
continue;
}
state.update(&key.0)?;
aside += rows::footprint(&key.0);
kept.insert(key);
}
}
_ => {
return Err(Error::internal(
"two instances of one DISTINCT aggregate disagree about what their sets hold",
));
}
}
}
Ok(aside)
}
fn tables(
table: &Table,
states: &Vec<Accumulator>,
counts: &Vec<i64>,
compact: &Vec<CompactNumeric>,
overflow: &HashMap<usize, (i128, i128)>,
seen: &Vec<DistinctSet>,
) -> u64 {
let width = |count: usize, size: usize| {
u64::try_from(count).unwrap_or(u64::MAX).saturating_mul(width_of(size))
};
table.footprint()
+ width(states.capacity(), size_of::<Accumulator>())
+ width(counts.capacity(), size_of::<i64>())
+ width(compact.capacity(), size_of::<CompactNumeric>())
+ overflow_footprint(overflow)
+ width(seen.capacity(), size_of::<DistinctSet>())
}
fn overflow_footprint(overflow: &HashMap<usize, (i128, i128)>) -> u64 {
width_of(overflow.capacity() * size_of::<(usize, (i128, i128))>() * 2)
}
fn width_of(size: usize) -> u64 {
u64::try_from(size).unwrap_or(u64::MAX)
}
fn group_name(plan: &Plan, group: ExprRef, input: &Schema, at: usize) -> String {
if let Expr::Column(binding) = *plan.expr(group)
&& let Some(position) = input.position_of(binding)
{
return input.fields()[position].name.clone();
}
format!("group{at}")
}
#[derive(Debug)]
pub(crate) struct Distinct {
on: Prepared,
whole: bool,
types: Vec<LogicalType>,
memory: Memory,
global: Mutex<Held>,
charged: Mutex<Vec<Reservation>>,
held: Mutex<Reservation>,
out: Buffered,
}
#[derive(Debug)]
struct Held {
seen: RowSet,
kept: Vec<Vec<Value>>,
counted: u64,
table: Reservation,
counted_rows: u64,
slots: Reservation,
}
#[derive(Debug)]
pub(crate) struct Keeping {
seen: RowSet,
kept: Vec<(Key, Vec<Value>)>,
scratch: Scratch,
key: Key,
rows: Reservation,
table: Reservation,
counted: u64,
counted_table: u64,
}
impl Distinct {
#[must_use]
pub(crate) fn in_session(mut self, session: &Session) -> Self {
self.on = self.on.in_session(session);
self
}
pub(crate) fn new(
plan: &Plan,
input: &Schema,
on: Slice,
memory: &Memory,
) -> Result<(Self, Buffered)> {
let on = plan.expr_list(on).to_vec();
let out = Buffered::new();
let distinct = Self {
whole: on.is_empty(),
on: Prepared::new(plan, &on, input)?,
types: input.types(),
memory: memory.clone(),
global: Mutex::new(Held {
seen: RowSet::default(),
kept: Vec::new(),
counted: 0,
table: memory.reservation(),
counted_rows: 0,
slots: memory.reservation(),
}),
charged: Mutex::new(Vec::new()),
held: Mutex::new(memory.reservation()),
out: out.clone(),
};
Ok((distinct, out))
}
}
impl Sink for Distinct {
type Local = Keeping;
fn parallel(&self) -> bool {
self.whole
}
fn local(&self) -> Keeping {
Keeping {
seen: RowSet::default(),
kept: Vec::new(),
scratch: self.on.scratch(),
key: Key(Vec::new()),
rows: self.memory.reservation(),
table: self.memory.reservation(),
counted: 0,
counted_table: 0,
}
}
fn sink(&self, chunk: &Chunk, local: &mut Keeping) -> Result<Progress> {
let mut keys = Vec::with_capacity(self.on.len());
self.on.evaluate(chunk, &mut local.scratch, &mut keys)?;
let mut taken = 0;
let mut aside = 0;
for row in 0..chunk.len() {
if self.whole {
local.key.0.clear();
local.key.0 = (0..chunk.width())
.map(|column| chunk.try_value_at(row, column))
.collect::<Result<_>>()?;
} else {
fill(&mut local.key, &keys, row)?;
}
if local.seen.contains(&local.key) {
continue;
}
let values: Vec<Value> = if self.whole {
local.key.0.clone()
} else {
(0..chunk.width())
.map(|column| chunk.try_value_at(row, column))
.collect::<Result<_>>()?
};
let stored = local.key.clone();
taken += rows::heap(&values) + rows::heap(&stored.0);
aside += rows::heap(&stored.0);
local.seen.insert(stored.clone());
local.kept.push((stored, values));
}
local.rows.grow(taken)?;
local.table.grow(aside)?;
let held = width_of(local.kept.capacity() * size_of::<(Key, Vec<Value>)>());
rows::capacity(held, &mut local.counted, &mut local.rows)?;
let now = rows::buckets(local.seen.capacity()) * (width_of(size_of::<Key>()) + 1);
rows::capacity(now, &mut local.counted_table, &mut local.table)?;
Ok(Progress::More)
}
fn combine(&self, local: Keeping) -> Result<()> {
let Keeping { seen, kept, rows, mut table, .. } = local;
drop(seen);
table.release();
let mut global = self.global.lock().map_err(poisoned)?;
let global = &mut *global;
let mut aside = 0;
for (key, values) in kept {
let cost = rows::heap(&key.0);
if global.seen.insert(key) {
aside += cost;
global.kept.push(values);
}
}
global.table.grow(aside)?;
let now = rows::buckets(global.seen.capacity()) * (width_of(size_of::<Key>()) + 1);
rows::capacity(now, &mut global.counted, &mut global.table)?;
let held = width_of(global.kept.capacity() * size_of::<Vec<Value>>());
rows::capacity(held, &mut global.counted_rows, &mut global.slots)?;
self.charged.lock().map_err(poisoned)?.push(rows);
Ok(())
}
fn finalize(&self, _threads: &Lease<'_>) -> Result<()> {
let mut global = self.global.lock().map_err(poisoned)?;
let kept = std::mem::take(&mut global.kept);
global.seen = RowSet::default();
global.counted = 0;
global.table.release();
let mut held = self.held.lock().map_err(poisoned)?;
let chunks = rows::chunks(&self.types, &kept, &mut held)?;
self.out.fill(chunks)?;
global.counted_rows = 0;
global.slots.release();
self.charged.lock().map_err(poisoned)?.clear();
Ok(())
}
}
fn poisoned<T>(_: T) -> Error {
Error::internal("a thread panicked while holding the rows a distinct is keeping")
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use std::sync::Arc;
use rudb_common::{Field, LogicalType, Memory, Value};
use rudb_kernels::NOWHERE;
use rudb_pipeline::Sink;
use rudb_plan::{Plan, Slice};
use rudb_vector::{Chunk, Data, Vector};
use super::{
Aggregate, BigIntDistinct, BigIntDistinctRuns, COMPACT_FROM, Call, CompactNumeric,
Distinct, EncodedCountRecord, EncodedCountRun, EncodedCountRuns, EncodedValid,
FixedPartition, FixedRecord, FixedRun, FixedRuns, PARTITION_FROM, RADIX_PARTITIONS,
RUN_BLOCK, Second, Share, Signed, Tucked, WINDOW_RATE, WINDOW_SLACK,
bigint_distinct_partition, encoded_count_partition, fixed_partition, interior,
slot_runs_of, spread_runs, spread_slots,
};
use crate::buffer::Buffered;
use crate::schema::Schema;
fn chunk(values: &[i32]) -> Chunk {
let column = Vector::flat(LogicalType::Integer, Data::Int32(values.to_vec().into()))
.expect("integers are an i32 layout");
Chunk::new(vec![column]).expect("one column is one length")
}
fn distinct() -> (Distinct, Buffered) {
let schema = Schema::numbered(vec![Field::new("a", LogicalType::Integer)], 0);
Distinct::new(&Plan::new(), &schema, Slice::EMPTY, &Memory::unlimited())
.expect("there are no expressions to resolve")
}
fn column(out: &Buffered) -> Vec<Value> {
let chunk = out.at(0).expect("readable").expect("one chunk");
(0..chunk.len()).map(|row| chunk.value_at(row, 0)).collect()
}
#[test]
fn runs_of_kept_rows_spread_like_their_slots() {
let all = 200;
let kept: Vec<u32> =
(0..all as u32).filter(|row| !matches!(row % 37, 0 | 5 | 6) && *row < 190).collect();
let slots: Vec<usize> = (0..kept.len()).map(|at| [3, 3, 9, 1][at / 30 % 4]).collect();
let mut runs = Vec::new();
assert!(slot_runs_of(&slots, &mut runs, 1));
let mut spread = Vec::new();
assert!(spread_runs(&runs, &kept, all, usize::MAX, &mut spread));
let mut moved = slots.clone();
spread_slots(&mut moved, &kept, all);
let mut expected = Vec::new();
assert!(slot_runs_of(&moved, &mut expected, usize::MAX));
assert_eq!(spread, expected);
assert!(!spread_runs(&runs, &kept, all, 3, &mut spread));
assert!(spread.is_empty());
let whole: Vec<u32> = (0..all as u32).collect();
assert!(spread_runs(&[(5, 150), (2, 200)], &whole, all, 2, &mut spread));
assert_eq!(spread, [(5, 150), (2, 200)]);
}
#[test]
fn slots_in_runs_are_cut_into_them_and_slots_in_no_order_are_not() {
let lengths = [(4, 40), (NOWHERE, 17), (1, 1), (4, 30), (0, 16)];
let slots: Vec<usize> =
lengths.iter().flat_map(|&(slot, length)| std::iter::repeat_n(slot, length)).collect();
let mut runs = Vec::new();
assert!(slot_runs_of(&slots, &mut runs, 1));
let mut end = 0;
let expected: Vec<(usize, usize)> = lengths
.iter()
.map(|&(slot, length)| {
end += length;
(slot, end)
})
.collect();
assert_eq!(runs, expected);
let one = vec![7; 1_000];
assert!(slot_runs_of(&one, &mut runs, 1));
assert_eq!(runs, [(7, 1_000)]);
let scattered: Vec<usize> = (0..1_000).map(|row| row * 7 % 13).collect();
assert!(!slot_runs_of(&scattered, &mut runs, 1));
assert!(runs.is_empty());
assert!(!slot_runs_of(&[], &mut runs, 1));
}
#[test]
fn slots_too_broken_up_for_one_call_are_cut_for_eight_of_them() {
let scattered: Vec<usize> = (0..1_000).map(|row| row * 7 % 13).collect();
let mut runs = Vec::new();
assert!(!slot_runs_of(&scattered, &mut runs, 1));
assert!(slot_runs_of(&scattered, &mut runs, 8));
assert_eq!(runs.len(), 1_000);
assert_eq!(runs.last(), Some(&(scattered[999], 1_000)));
assert!(!slot_runs_of(&scattered, &mut runs, 0));
assert!(runs.is_empty());
}
#[test]
fn runs_ending_on_the_edge_of_a_block_are_cut_where_they_end() {
for first in [RUN_BLOCK - 1, RUN_BLOCK, RUN_BLOCK + 1] {
let lengths = [(3, first), (8, 1), (3, RUN_BLOCK * 2), (5, 2)];
let slots: Vec<usize> = lengths
.iter()
.flat_map(|&(slot, length)| std::iter::repeat_n(slot, length))
.collect();
let mut runs = Vec::new();
assert!(slot_runs_of(&slots, &mut runs, 8), "gave up on runs ending at {first}");
let mut end = 0;
let expected: Vec<(usize, usize)> = lengths
.iter()
.map(|&(slot, length)| {
end += length;
(slot, end)
})
.collect();
assert_eq!(runs, expected, "the runs ending at {first} were cut wrong");
}
}
#[test]
fn one_instance_keeps_the_first_of_each_row() {
let (distinct, out) = distinct();
let mut local = distinct.local();
distinct.sink(&chunk(&[1, 2, 1, 3, 2]), &mut local).expect("five rows");
distinct.combine(local).expect("the one instance");
distinct.finalize(&rudb_pipeline::Lease::alone()).expect("the answer");
assert_eq!(column(&out), [Value::Integer(1), Value::Integer(2), Value::Integer(3)]);
}
#[test]
fn two_instances_that_both_kept_a_row_keep_one_of_it_between_them() {
let (distinct, out) = distinct();
let mut left = distinct.local();
let mut right = distinct.local();
distinct.sink(&chunk(&[1, 2]), &mut left).expect("two rows");
distinct.sink(&chunk(&[2, 3]), &mut right).expect("two rows");
distinct.combine(left).expect("the first instance");
distinct.combine(right).expect("the second instance");
distinct.finalize(&rudb_pipeline::Lease::alone()).expect("the answer");
assert_eq!(column(&out), [Value::Integer(1), Value::Integer(2), Value::Integer(3)]);
}
#[test]
fn an_ungrouped_aggregate_answers_one_row_from_one_instance() {
let plan = Plan::new();
let schema = Schema::numbered(vec![Field::new("a", LogicalType::Integer)], 0);
let (aggregate, out) =
Aggregate::new(&plan, &schema, 1, Slice::EMPTY, Slice::EMPTY, &Memory::unlimited())
.expect("no aggregates to take apart");
aggregate.combine(aggregate.local()).expect("the one instance");
aggregate.finalize(&rudb_pipeline::Lease::alone()).expect("the answer");
assert_eq!(out.at(0).expect("readable").expect("one chunk").len(), 1);
}
fn parsed(text: &str) -> Plan {
Plan::parse(&format!("{text}\n Get memory.main.t AS t #0 [x::INTEGER]"))
.expect("a plan this crate's own notation describes")
}
fn aggregate(plan: &Plan) -> (Aggregate<'_>, Buffered) {
let schema = Schema::numbered(vec![Field::new("x", LogicalType::Integer)], 0);
let (groups, aggregates) = match *plan.node(plan.root()) {
rudb_plan::Node::Aggregate { groups, aggregates, .. } => (groups, aggregates),
ref other => panic!("the plan's root is {other:?} and not an aggregate"),
};
Aggregate::new(plan, &schema, 1, groups, aggregates, &Memory::unlimited())
.expect("the aggregates are ones this crate implements")
}
fn answer(out: &Buffered) -> Vec<Vec<Value>> {
let mut rows: Vec<Vec<Value>> = Vec::new();
for at in 0.. {
let Some(chunk) = out.at(at).expect("readable") else { break };
for row in 0..chunk.len() {
rows.push((0..chunk.width()).map(|column| chunk.value_at(row, column)).collect());
}
}
rows
}
fn bigint_chunk(values: &[Value]) -> Chunk {
let column = Vector::from_values(LogicalType::BigInt, values).expect("BIGINT values");
Chunk::new(vec![column]).expect("one column is one length")
}
#[test]
fn radix_bigint_distinct_counts_across_instances_and_skips_nulls() {
let plan = Plan::parse(concat!(
"Aggregate #1 groups=[] aggregates=[count(DISTINCT #0.0::BIGINT)::BIGINT]\n",
" Get memory.main.t AS t #0 [x::BIGINT]",
))
.expect("a distinct count plan");
let schema = Schema::numbered(vec![Field::new("x", LogicalType::BigInt)], 0);
let rudb_plan::Node::Aggregate { groups, aggregates, .. } = *plan.node(plan.root()) else {
panic!("the root is an aggregate")
};
let (aggregate, out) =
Aggregate::new(&plan, &schema, 1, groups, aggregates, &Memory::unlimited())
.expect("a distinct count aggregate");
let mut left = aggregate.local();
let mut right = aggregate.local();
aggregate
.sink(&bigint_chunk(&[Value::BigInt(7), Value::Null, Value::BigInt(8)]), &mut left)
.expect("the left values");
aggregate
.sink(&bigint_chunk(&[Value::BigInt(8), Value::BigInt(9), Value::Null]), &mut right)
.expect("the right values");
aggregate.combine(left).expect("the left instance");
aggregate.combine(right).expect("the right instance");
aggregate.finalize(&rudb_pipeline::Lease::alone()).expect("the distinct count");
assert_eq!(answer(&out), [vec![Value::BigInt(3)]]);
}
#[test]
fn radix_bigint_distinct_answers_zero_without_input() {
let plan = Plan::parse(concat!(
"Aggregate #1 groups=[] aggregates=[count(DISTINCT #0.0::BIGINT)::BIGINT]\n",
" Get memory.main.t AS t #0 [x::BIGINT]",
))
.expect("a distinct count plan");
let schema = Schema::numbered(vec![Field::new("x", LogicalType::BigInt)], 0);
let rudb_plan::Node::Aggregate { groups, aggregates, .. } = *plan.node(plan.root()) else {
panic!("the root is an aggregate")
};
let (aggregate, out) =
Aggregate::new(&plan, &schema, 1, groups, aggregates, &Memory::unlimited())
.expect("a distinct count aggregate");
aggregate.combine(aggregate.local()).expect("an empty instance");
aggregate.finalize(&rudb_pipeline::Lease::alone()).expect("the zero count");
assert_eq!(answer(&out), [vec![Value::BigInt(0)]]);
}
#[test]
fn two_instances_of_a_grouped_aggregate_answer_one_row_a_group() {
let plan = parsed("Aggregate #1 groups=[#0.0::INTEGER] aggregates=[count_star()::BIGINT]");
let (aggregate, out) = aggregate(&plan);
let mut left = aggregate.local();
let mut right = aggregate.local();
aggregate.sink(&chunk(&[1, 2, 1]), &mut left).expect("three rows");
aggregate.sink(&chunk(&[2, 3, 2]), &mut right).expect("three rows");
aggregate.combine(left).expect("the first instance");
aggregate.combine(right).expect("the second instance");
aggregate.finalize(&rudb_pipeline::Lease::alone()).expect("the answer");
let mut rows = answer(&out);
rows.sort_by_key(|row| format!("{:?}", row[0]));
assert_eq!(
rows,
[
vec![Value::Integer(1), Value::BigInt(2)],
vec![Value::Integer(2), Value::BigInt(3)],
vec![Value::Integer(3), Value::BigInt(1)],
]
);
}
#[test]
fn two_instances_under_a_limit_keep_the_same_groups() {
let plan = parsed("Aggregate #1 groups=[#0.0::INTEGER] aggregates=[count_star()::BIGINT]");
let (aggregate, out) = aggregate(&plan);
let aggregate = aggregate.limit_groups(2);
let mut left = aggregate.local();
let mut right = aggregate.local();
aggregate.sink(&chunk(&[1, 2, 1, 2]), &mut left).expect("the left rows");
aggregate.sink(&chunk(&[3, 4, 1, 2]), &mut right).expect("the right rows");
aggregate.combine(left).expect("the left instance");
aggregate.combine(right).expect("the right instance");
aggregate.finalize(&rudb_pipeline::Lease::alone()).expect("the answer");
let mut rows = answer(&out);
rows.sort_by_key(|row| format!("{:?}", row[0]));
assert_eq!(
rows,
[vec![Value::Integer(1), Value::BigInt(3)], vec![Value::Integer(2), Value::BigInt(3)],]
);
}
#[test]
fn two_instances_under_a_limit_nothing_reaches_keep_every_group() {
let plan = parsed("Aggregate #1 groups=[#0.0::INTEGER] aggregates=[count_star()::BIGINT]");
let (aggregate, out) = aggregate(&plan);
let aggregate = aggregate.limit_groups(10);
let mut left = aggregate.local();
let mut right = aggregate.local();
aggregate.sink(&chunk(&[1, 2]), &mut left).expect("the left rows");
aggregate.sink(&chunk(&[2, 3]), &mut right).expect("the right rows");
aggregate.combine(left).expect("the left instance");
aggregate.combine(right).expect("the right instance");
aggregate.finalize(&rudb_pipeline::Lease::alone()).expect("the answer");
let mut rows = answer(&out);
rows.sort_by_key(|row| format!("{:?}", row[0]));
assert_eq!(
rows,
[
vec![Value::Integer(1), Value::BigInt(1)],
vec![Value::Integer(2), Value::BigInt(2)],
vec![Value::Integer(3), Value::BigInt(1)],
]
);
}
#[test]
fn a_partitioned_aggregate_answers_every_group_once() {
let plan = parsed("Aggregate #1 groups=[#0.0::INTEGER] aggregates=[count_star()::BIGINT]");
let (aggregate, out) = aggregate(&plan);
let mut left = aggregate.local();
let mut right = aggregate.local();
let values: Vec<i32> = (0..20_000).collect();
for part in values.chunks(1_024) {
aggregate.sink(&chunk(part), &mut left).expect("a chunk of groups");
aggregate.sink(&chunk(part), &mut right).expect("the same groups again");
}
aggregate.combine(left).expect("the first instance");
aggregate.combine(right).expect("the second instance");
let built = aggregate.built.lock().expect("readable");
assert!(
built.partitioning,
"twenty thousand groups on two instances is meant to take the partitioned path"
);
assert!(built.local, "twenty thousand groups fit in the cache twice over");
drop(built);
aggregate.finalize(&rudb_pipeline::Lease::alone()).expect("the answer");
let mut seen: Vec<i32> = Vec::new();
for row in answer(&out) {
assert_eq!(row[1], Value::BigInt(2), "{row:?} was counted on both instances");
match row[0] {
Value::Integer(key) => seen.push(key),
ref other => panic!("the group is {other:?} and not an integer"),
}
}
seen.sort_unstable();
assert_eq!(seen, values);
}
#[test]
fn an_aggregate_too_large_for_the_cache_gives_up_its_own_tables_and_still_answers_once() {
let plan = parsed(concat!(
"Aggregate #1 groups=[#0.0::INTEGER] aggregates=[count_star()::BIGINT, ",
"sum(#0.0::INTEGER)::BIGINT, min(#0.0::INTEGER)::INTEGER]"
));
let (aggregate, out) = aggregate(&plan);
let mut left = aggregate.local();
let mut right = aggregate.local();
let values: Vec<i32> = (0..120_000).collect();
for part in values.chunks(1_024) {
for instance in [&mut left, &mut right] {
aggregate.sink(&chunk(part), instance).expect("a chunk of groups");
aggregate.sink(&chunk(part), instance).expect("the same chunk a second time");
}
}
assert!(
!aggregate.built.lock().expect("readable").local,
"a hundred and twenty thousand groups of three calls is past what the cache holds"
);
aggregate.combine(left).expect("the first instance");
aggregate.combine(right).expect("the second instance");
aggregate.finalize(&rudb_pipeline::Lease::alone()).expect("the answer");
let mut seen: Vec<i32> = Vec::new();
for row in answer(&out) {
let Value::Integer(key) = row[0] else {
panic!("the group is {:?} and not an integer", row[0]);
};
assert_eq!(row[1], Value::BigInt(4), "{row:?} was not counted four times");
assert_eq!(
row[2],
Value::BigInt(i64::from(key) * 4),
"{row:?} was not summed four times"
);
assert_eq!(row[3], Value::Integer(key), "{row:?} kept the wrong least value");
seen.push(key);
}
seen.sort_unstable();
assert_eq!(seen, values);
}
#[test]
fn an_aggregate_whose_keys_arrive_in_runs_keeps_its_own_tables() {
let plan = parsed(concat!(
"Aggregate #1 groups=[#0.0::INTEGER] aggregates=[count_star()::BIGINT, ",
"sum(#0.0::INTEGER)::BIGINT, min(#0.0::INTEGER)::INTEGER]"
));
let (aggregate, out) = aggregate(&plan);
let mut left = aggregate.local();
let mut right = aggregate.local();
let values: Vec<i32> = (0..120_000).collect();
let runs: Vec<i32> = values.iter().flat_map(|&key| [key, key]).collect();
for part in runs.chunks(1_024) {
for instance in [&mut left, &mut right] {
aggregate.sink(&chunk(part), instance).expect("a chunk of runs");
}
}
assert!(
aggregate.built.lock().expect("readable").local,
"keys that arrive in runs are already divided between the instances"
);
aggregate.combine(left).expect("the first instance");
aggregate.combine(right).expect("the second instance");
aggregate.finalize(&rudb_pipeline::Lease::alone()).expect("the answer");
let mut seen: Vec<i32> = Vec::new();
for row in answer(&out) {
let Value::Integer(key) = row[0] else {
panic!("the group is {:?} and not an integer", row[0]);
};
assert_eq!(row[1], Value::BigInt(4), "{row:?} was not counted four times");
seen.push(key);
}
seen.sort_unstable();
assert_eq!(seen, values);
}
#[test]
fn an_aggregate_under_a_pushed_down_bound_is_split_all_the_same() {
let plan = parsed(
"Aggregate #1 groups=[#0.0::INTEGER] aggregates=[count_star()::BIGINT, count_star()::BIGINT]",
);
let (aggregate, out) = aggregate(&plan);
let aggregate = aggregate.top_counts(1_000, 0);
let mut left = aggregate.local();
let mut right = aggregate.local();
let values: Vec<i32> = (0..20_000).collect();
for part in values.chunks(1_024) {
aggregate.sink(&chunk(part), &mut left).expect("a chunk of groups");
aggregate.sink(&chunk(part), &mut right).expect("the same groups again");
}
aggregate.combine(left).expect("the first instance");
aggregate.combine(right).expect("the second instance");
assert!(
aggregate.built.lock().expect("readable").partitioning,
"twenty thousand groups is over the line whatever the bound above says"
);
aggregate.finalize(&rudb_pipeline::Lease::alone()).expect("the answer");
assert_eq!(
answer(&out).len(),
20_000,
"a bound over what a partition holds throws nothing away"
);
}
#[test]
fn a_counted_exchange_adds_instances_and_keeps_the_largest_groups() {
let plan = parsed("Aggregate #1 groups=[#0.0::INTEGER] aggregates=[count_star()::BIGINT]");
let (aggregate, out) = aggregate(&plan);
let aggregate = aggregate.top_counts(2, 0);
assert!(aggregate.counted_top_count());
let mut left = aggregate.local();
let mut right = aggregate.local();
let values: Vec<i32> = (0..50_000).collect();
for part in values.chunks(1_024) {
aggregate.sink(&chunk(part), &mut left).expect("a chunk of groups");
}
aggregate.sink(&chunk(&[7, 7, 7, 9, 7]), &mut right).expect("a run and a repeat");
let nulls = Vector::from_values(LogicalType::Integer, &[Value::Null, Value::Integer(9)])
.expect("INTEGER values");
aggregate
.sink(&Chunk::new(vec![nulls]).expect("one column"), &mut right)
.expect("a null key");
aggregate.combine(left).expect("the first instance");
aggregate.combine(right).expect("the second instance");
aggregate.finalize(&rudb_pipeline::Lease::alone()).expect("the answer");
let mut rows = answer(&out);
rows.sort_by_key(|row| match row[1] {
Value::BigInt(count) => std::cmp::Reverse(count),
ref other => panic!("a count of {other:?}"),
});
assert_eq!(rows[0], vec![Value::Integer(7), Value::BigInt(5)]);
assert_eq!(rows[1], vec![Value::Integer(9), Value::BigInt(3)]);
assert!(rows.contains(&vec![Value::Null, Value::BigInt(1)]), "the null group is kept");
assert!(rows.len() < 50_000, "each split keeps its two largest and no more");
}
#[test]
fn two_instances_of_an_ungrouped_aggregate_add_up_to_one_total() {
let plan = parsed(
"Aggregate #1 groups=[] aggregates=[sum(#0.0::INTEGER)::HUGEINT, min(#0.0::INTEGER)::INTEGER, max(#0.0::INTEGER)::INTEGER, count_star()::BIGINT]",
);
let (aggregate, out) = aggregate(&plan);
let mut left = aggregate.local();
let mut right = aggregate.local();
aggregate.sink(&chunk(&[4, 7]), &mut left).expect("two rows");
aggregate.sink(&chunk(&[2, 9]), &mut right).expect("two rows");
aggregate.combine(left).expect("the first instance");
aggregate.combine(right).expect("the second instance");
aggregate.finalize(&rudb_pipeline::Lease::alone()).expect("the answer");
assert_eq!(
answer(&out),
[vec![Value::HugeInt(22), Value::Integer(2), Value::Integer(9), Value::BigInt(4)]]
);
}
#[test]
fn an_instance_that_saw_no_rows_changes_nothing_when_it_combines() {
let plan = parsed("Aggregate #1 groups=[#0.0::INTEGER] aggregates=[count_star()::BIGINT]");
let (aggregate, out) = aggregate(&plan);
let mut seen = aggregate.local();
aggregate.sink(&chunk(&[5, 5]), &mut seen).expect("two rows");
aggregate.combine(aggregate.local()).expect("an instance that saw nothing");
aggregate.combine(seen).expect("the one that saw something");
aggregate.combine(aggregate.local()).expect("another that saw nothing");
aggregate.finalize(&rudb_pipeline::Lease::alone()).expect("the answer");
assert_eq!(answer(&out), [vec![Value::Integer(5), Value::BigInt(2)]]);
}
#[test]
fn two_instances_of_a_distinct_aggregate_count_a_shared_value_once() {
let plan = parsed(
"Aggregate #1 groups=[#0.0::INTEGER] aggregates=[count(DISTINCT #0.0::INTEGER)::BIGINT]",
);
let (aggregate, out) = aggregate(&plan);
let mut left = aggregate.local();
let mut right = aggregate.local();
aggregate.sink(&chunk(&[3, 3]), &mut left).expect("two rows of one group");
aggregate.sink(&chunk(&[3, 4]), &mut right).expect("the same group and another");
aggregate.combine(left).expect("the first instance");
aggregate.combine(right).expect("the second instance");
aggregate.finalize(&rudb_pipeline::Lease::alone()).expect("the answer");
let mut rows = answer(&out);
rows.sort_by_key(|row| format!("{:?}", row[0]));
assert_eq!(
rows,
[vec![Value::Integer(3), Value::BigInt(1)], vec![Value::Integer(4), Value::BigInt(1)]]
);
}
#[test]
fn a_distinct_over_nothing_produces_nothing() {
let (distinct, out) = distinct();
distinct.combine(distinct.local()).expect("an instance that saw no chunks");
distinct.finalize(&rudb_pipeline::Lease::alone()).expect("the answer");
assert_eq!(out.len().expect("readable"), 0);
}
#[test]
fn compact_smallint_totals_remain_exact_past_i64() {
let mut wide = CompactNumeric { sum: i64::MAX, mean: i64::MIN, ..Default::default() };
let mut wide_overflow = HashMap::new();
wide.add(0, Some(1), Some(-1), &mut wide_overflow).expect("wide totals");
assert_eq!(
wide.totals(0, &wide_overflow),
(i128::from(i64::MAX) + 1, i128::from(i64::MIN) - 1)
);
let mut coming = CompactNumeric::default();
let mut coming_overflow = HashMap::new();
coming.add(0, Some(2), Some(3), &mut coming_overflow).expect("small totals");
wide.combine(0, &coming, 0, &coming_overflow, &mut wide_overflow).expect("combined totals");
assert_eq!(
wide.totals(0, &wide_overflow),
(i128::from(i64::MAX) + 3, i128::from(i64::MIN) + 2)
);
assert_eq!(wide.count(), 2);
assert_eq!(wide.mean_count, 2);
assert_eq!(size_of::<CompactNumeric>(), 32);
}
#[test]
fn a_group_that_fits_stays_out_of_a_map_another_group_has_already_used() {
let mut overflow = HashMap::new();
let mut wide = CompactNumeric { sum: i64::MAX, ..Default::default() };
wide.add(4, Some(1), None, &mut overflow).expect("wide totals");
assert_eq!(overflow.len(), 1, "the wide group is the only one in the map");
let mut narrow = CompactNumeric { sum: 10, ..Default::default() };
narrow.add(9, Some(5), None, &mut overflow).expect("small totals");
assert_eq!(overflow.len(), 1, "a group that fits sixty four bits is not written down");
assert_eq!(narrow.totals(9, &overflow), (15, 0));
assert_eq!(wide.totals(4, &overflow), (i128::from(i64::MAX) + 1, 0));
let zero = CompactNumeric::default();
let mut coming = HashMap::new();
wide.combine(4, &zero, 0, &coming, &mut overflow).expect("no change to the total");
assert_eq!(wide.totals(4, &overflow), (i128::from(i64::MAX) + 1, 0));
let mut back = CompactNumeric { sum: i64::MIN + 1, ..Default::default() };
back.add(0, None, None, &mut coming).expect("a total that fits");
wide.combine(4, &back, 0, &coming, &mut overflow).expect("back under the limit");
assert!(overflow.is_empty(), "the group left the map when its total fitted again");
assert_eq!(wide.totals(4, &overflow), (i128::from(i64::MAX) + i128::from(i64::MIN) + 2, 0));
}
#[test]
fn a_bigint_distinct_set_allocates_only_after_its_first_value() {
let mut values = BigIntDistinct::default();
assert!(values.insert(7));
assert!(!values.insert(7));
assert!(matches!(values, BigIntDistinct::One(7)));
assert!(values.insert(9));
assert!(!values.insert(7));
assert!(!values.insert(9));
assert!(matches!(values, BigIntDistinct::Many(_)));
}
#[test]
fn a_bigint_radix_partition_counts_unique_values_across_the_runs_it_was_handed() {
let mut partition =
BigIntDistinctRuns { runs: vec![vec![11, 12, 11], vec![13, 11], Vec::new(), vec![12]] };
assert_eq!(
bigint_distinct_partition(&mut partition, &Memory::unlimited())
.expect("the distinct partition"),
3,
"a value counts once however many instances handed it over"
);
}
#[test]
fn a_bigint_radix_partition_tells_apart_values_that_probe_past_each_other() {
let mut partition = BigIntDistinctRuns {
runs: vec![(0..300).map(i64::from).collect(), (150..450).map(i64::from).collect()],
};
assert_eq!(
bigint_distinct_partition(&mut partition, &Memory::unlimited())
.expect("the distinct partition"),
450,
"four hundred and fifty different values are four hundred and fifty groups"
);
}
#[test]
fn a_two_key_encoded_count_record_holds_no_second_key() {
assert_eq!(size_of::<EncodedCountRecord<()>>(), 16);
assert_eq!(size_of::<EncodedCountRecord<Tucked>>(), 16);
assert_eq!(size_of::<EncodedCountRecord<i64>>(), 24);
}
#[test]
fn an_encoded_count_partition_aggregates_collisions_and_nulls_exactly() {
let dictionary = Vector::from_values(
LogicalType::Varchar,
&[Value::Varchar("one".into()), Value::Varchar("two".into())],
)
.expect("a string dictionary");
let row = |first, second, third| EncodedCountRecord { first, second, hash: 7, third };
let mut partition = EncodedCountRun::default();
partition.push(row(1, 2, 0), EncodedValid::ALL);
partition.push(row(1, 2, 0), EncodedValid::ALL);
partition.push(row(1, 2, 1), EncodedValid::ALL);
partition.push(row(0, 2, 0), EncodedValid::SECOND | EncodedValid::THIRD);
let leading = [LogicalType::BigInt, LogicalType::BigInt];
let part = encoded_count_partition(
&mut EncodedCountRuns { runs: vec![partition] },
&dictionary,
&leading,
u32::BITS,
10,
&Memory::unlimited(),
)
.expect("the encoded partition");
let mut rows: Vec<Vec<Value>> = Vec::new();
for chunk in part.chunks {
for row in 0..chunk.len() {
rows.push((0..chunk.width()).map(|column| chunk.value_at(row, column)).collect());
}
}
rows.sort_by_key(|row| format!("{row:?}"));
let mut expected = vec![
vec![
Value::BigInt(1),
Value::BigInt(2),
Value::Varchar("one".into()),
Value::BigInt(2),
],
vec![
Value::BigInt(1),
Value::BigInt(2),
Value::Varchar("two".into()),
Value::BigInt(1),
],
vec![Value::Null, Value::BigInt(2), Value::Varchar("one".into()), Value::BigInt(1)],
];
expected.sort_by_key(|row| format!("{row:?}"));
assert_eq!(rows, expected);
assert_eq!(size_of::<EncodedCountRecord>(), 24);
}
#[test]
fn an_encoded_count_partition_split_for_cache_adds_weights_across_runs_and_splits() {
let dictionary = Vector::from_values(LogicalType::Varchar, &[Value::Varchar("one".into())])
.expect("a string dictionary");
let groups = (super::FIXED_SPLIT_ROWS * 8) as i64;
let row = |first: i64| EncodedCountRecord {
first,
second: 0,
hash: (first as u32).wrapping_mul(0x9e37_79b9),
third: 0,
};
let mut early = EncodedCountRun::default();
let mut late = EncodedCountRun::default();
for first in 0..groups {
let times = if first % 10_000 == 0 { 2 + first / 10_000 } else { 1 };
for time in 0..times {
let into = if time % 2 == 0 { &mut early } else { &mut late };
into.push(row(first), EncodedValid::ALL);
}
}
early.compact();
let part = encoded_count_partition(
&mut EncodedCountRuns { runs: vec![early, late] },
&dictionary,
&[LogicalType::BigInt],
u32::BITS,
3,
&Memory::unlimited(),
)
.expect("the encoded partition");
let mut rows: Vec<Vec<Value>> = Vec::new();
for chunk in part.chunks {
for row in 0..chunk.len() {
rows.push((0..chunk.width()).map(|column| chunk.value_at(row, column)).collect());
}
}
let largest = groups / 10_000 * 10_000;
let expected = (0..3)
.map(|step| {
let first = largest - step * 10_000;
vec![
Value::BigInt(first),
Value::Varchar("one".into()),
Value::BigInt(2 + first / 10_000),
]
})
.collect::<Vec<_>>();
assert_eq!(rows, expected, "the three largest groups, largest first");
}
#[test]
fn an_encoded_count_partition_folds_every_run_into_the_widest_one() {
let dictionary = Vector::from_values(
LogicalType::Varchar,
&[Value::Varchar("one".into()), Value::Varchar("two".into())],
)
.expect("a string dictionary");
let row = |first, third| EncodedCountRecord { first, second: 0, hash: 7, third };
let mut narrow = EncodedCountRun::default();
narrow.push(row(1, 0), EncodedValid::ALL);
let mut widest = EncodedCountRun::default();
for _ in 0..3 {
widest.push(row(1, 0), EncodedValid::ALL);
}
widest.push(row(2, 1), EncodedValid::ALL);
let mut late = EncodedCountRun::default();
late.push(row(0, 0), EncodedValid::SECOND | EncodedValid::THIRD);
late.push(row(1, 0), EncodedValid::ALL);
let leading = [LogicalType::BigInt];
let part = encoded_count_partition(
&mut EncodedCountRuns { runs: vec![narrow, widest, late] },
&dictionary,
&leading,
u32::BITS,
10,
&Memory::unlimited(),
)
.expect("the encoded partition");
let mut rows: Vec<Vec<Value>> = Vec::new();
for chunk in part.chunks {
for row in 0..chunk.len() {
rows.push((0..chunk.width()).map(|column| chunk.value_at(row, column)).collect());
}
}
rows.sort_by_key(|row| format!("{row:?}"));
let mut expected = vec![
vec![Value::BigInt(1), Value::Varchar("one".into()), Value::BigInt(5)],
vec![Value::BigInt(2), Value::Varchar("two".into()), Value::BigInt(1)],
vec![Value::Null, Value::Varchar("one".into()), Value::BigInt(1)],
];
expected.sort_by_key(|row| format!("{row:?}"));
assert_eq!(rows, expected, "a group is one group however many runs it arrived in");
}
#[test]
fn a_full_encoded_run_folds_its_repeats_and_the_counts_come_out_the_same() {
let dictionary = Vector::from_values(
LogicalType::Varchar,
&[Value::Varchar("one".into()), Value::Varchar("two".into())],
)
.expect("a string dictionary");
let row = |first, third| EncodedCountRecord { first, second: 0, hash: 7, third };
let mut folded = EncodedCountRun::default();
let rounds = COMPACT_FROM * 3;
for round in 0..rounds {
folded.scatter(row(round as i64 % 4, (round % 2) as u32), EncodedValid::ALL, 1);
if round % 3 == 0 {
folded.scatter(row(0, 1), EncodedValid::SECOND | EncodedValid::THIRD, 2);
}
}
assert!(
folded.len() <= COMPACT_FROM * 2 && folded.room <= COMPACT_FROM * 2,
"a run of five groups grew to {}",
folded.len()
);
let mut other = EncodedCountRun::default();
other.scatter(row(1, 1), EncodedValid::ALL, 5);
let leading = [LogicalType::BigInt];
let part = encoded_count_partition(
&mut EncodedCountRuns { runs: vec![other, folded] },
&dictionary,
&leading,
u32::BITS,
10,
&Memory::unlimited(),
)
.expect("the encoded partition");
let mut rows: Vec<Vec<Value>> = Vec::new();
for chunk in part.chunks {
for row in 0..chunk.len() {
rows.push((0..chunk.width()).map(|column| chunk.value_at(row, column)).collect());
}
}
rows.sort_by_key(|row| format!("{row:?}"));
let quarter = (rounds / 4) as i64;
let mut expected = vec![
vec![Value::BigInt(0), Value::Varchar("one".into()), Value::BigInt(quarter)],
vec![Value::BigInt(1), Value::Varchar("two".into()), Value::BigInt(quarter + 5)],
vec![Value::BigInt(2), Value::Varchar("one".into()), Value::BigInt(quarter)],
vec![Value::BigInt(3), Value::Varchar("two".into()), Value::BigInt(quarter)],
vec![
Value::Null,
Value::Varchar("two".into()),
Value::BigInt(rounds.div_ceil(3) as i64 * 2),
],
];
expected.sort_by_key(|row| format!("{row:?}"));
assert_eq!(rows, expected, "a folded run counts every row it was given");
}
#[test]
fn a_two_key_encoded_count_omits_the_unused_integer_and_narrows_the_one_it_keeps() {
let dictionary = Vector::from_values(
LogicalType::Varchar,
&[Value::Varchar("one".into()), Value::Varchar("two".into())],
)
.expect("a string dictionary");
let row = |first, third| EncodedCountRecord { first, second: 0, hash: 7, third };
let mut partition = EncodedCountRun::default();
partition.push(row(1, 0), EncodedValid::ALL);
partition.push(row(1, 0), EncodedValid::ALL);
partition.push(row(1, 1), EncodedValid::ALL);
partition.push(row(0, 0), EncodedValid::SECOND | EncodedValid::THIRD);
let leading = [LogicalType::SmallInt];
let part = encoded_count_partition(
&mut EncodedCountRuns { runs: vec![partition] },
&dictionary,
&leading,
u32::BITS,
10,
&Memory::unlimited(),
)
.expect("the encoded partition");
let mut rows: Vec<Vec<Value>> = Vec::new();
for chunk in part.chunks {
for row in 0..chunk.len() {
rows.push((0..chunk.width()).map(|column| chunk.value_at(row, column)).collect());
}
}
rows.sort_by_key(|row| format!("{row:?}"));
let mut expected = vec![
vec![Value::SmallInt(1), Value::Varchar("one".into()), Value::BigInt(2)],
vec![Value::SmallInt(1), Value::Varchar("two".into()), Value::BigInt(1)],
vec![Value::Null, Value::Varchar("one".into()), Value::BigInt(1)],
];
expected.sort_by_key(|row| format!("{row:?}"));
assert_eq!(rows, expected);
}
#[test]
fn a_bigint_and_stable_string_top_count_takes_the_encoded_path() {
let plan = Plan::parse(concat!(
"Aggregate #1 groups=[#0.0::BIGINT, #0.1::VARCHAR] ",
"aggregates=[count_star()::BIGINT]\n",
" Get memory.main.t AS t #0 [x::BIGINT, y::VARCHAR]",
))
.expect("a two-key count plan");
let schema = Schema::numbered(
vec![Field::new("x", LogicalType::BigInt), Field::new("y", LogicalType::Varchar)],
0,
);
let rudb_plan::Node::Aggregate { groups, aggregates, .. } = *plan.node(plan.root()) else {
panic!("the root is an aggregate")
};
let (aggregate, out) =
Aggregate::new(&plan, &schema, 1, groups, aggregates, &Memory::unlimited())
.expect("a count aggregate");
let aggregate = aggregate.top_counts(10, 0);
let users = Vector::from_values(
LogicalType::BigInt,
&[Value::BigInt(7), Value::BigInt(7), Value::BigInt(8)],
)
.expect("user ids");
let dictionary = Arc::new(
Vector::from_values(
LogicalType::Varchar,
&[Value::Varchar("one".into()), Value::Varchar("two".into())],
)
.expect("search phrases"),
);
let phrases = Vector::stable_dictionary(vec![0, 0, 1], dictionary)
.expect("stable search phrase codes");
let input = Chunk::new(vec![users, phrases]).expect("two aligned columns");
let mut local = aggregate.local();
aggregate.sink(&input, &mut local).expect("three rows");
aggregate.combine(local).expect("the one instance");
assert!(aggregate.encoded_count.get().is_some(), "the compact path was selected");
aggregate.finalize(&rudb_pipeline::Lease::alone()).expect("the answer");
let mut rows = answer(&out);
rows.sort_by_key(|row| format!("{row:?}"));
let mut expected = vec![
vec![Value::BigInt(7), Value::Varchar("one".into()), Value::BigInt(2)],
vec![Value::BigInt(8), Value::Varchar("two".into()), Value::BigInt(1)],
];
expected.sort_by_key(|row| format!("{row:?}"));
assert_eq!(rows, expected);
}
#[test]
fn a_null_in_the_first_record_of_a_partition_is_kept() {
let mut encoded = EncodedCountRun::default();
let row = EncodedCountRecord { first: 0, second: 0, hash: 7, third: 0 };
let some = EncodedValid::SECOND | EncodedValid::THIRD;
encoded.push(row, some);
encoded.push(row, EncodedValid::ALL);
assert_eq!(encoded.pending_validity, vec![some, EncodedValid::ALL]);
let mut fixed = FixedPartition::default();
let row = FixedRecord { first: 0, second: 0, sum: 0, mean: 0 };
let some = FixedRecord::SECOND | FixedRecord::SUM | FixedRecord::MEAN;
fixed.push(row, some);
fixed.push(row, FixedRecord::ALL);
assert_eq!(fixed.validity, vec![some, FixedRecord::ALL]);
}
fn two_key_counts(chunks: Vec<Chunk>) -> Vec<Vec<Value>> {
let plan = Plan::parse(concat!(
"Aggregate #1 groups=[#0.0::VARCHAR, #0.1::VARCHAR] ",
"aggregates=[count_star()::BIGINT]\n",
" Get memory.main.t AS t #0 [a::VARCHAR, b::VARCHAR]",
))
.expect("a two-key count plan");
let schema = Schema::numbered(
vec![Field::new("a", LogicalType::Varchar), Field::new("b", LogicalType::Varchar)],
0,
);
let rudb_plan::Node::Aggregate { groups, aggregates, .. } = *plan.node(plan.root()) else {
panic!("the root is an aggregate")
};
let (aggregate, out) =
Aggregate::new(&plan, &schema, 1, groups, aggregates, &Memory::unlimited())
.expect("a count aggregate");
let mut local = aggregate.local();
for chunk in &chunks {
aggregate.sink(chunk, &mut local).expect("a chunk of rows");
}
aggregate.combine(local).expect("the one instance");
aggregate.finalize(&rudb_pipeline::Lease::alone()).expect("the answer");
let mut rows = answer(&out);
rows.sort_by_key(|row| format!("{row:?}"));
rows
}
fn letters(values: &[Option<&str>]) -> Vector {
let values: Vec<Value> = values
.iter()
.map(|value| match value {
Some(text) => Value::Varchar((*text).into()),
None => Value::Null,
})
.collect();
Vector::from_values(LogicalType::Varchar, &values).expect("a column of strings")
}
fn letters_coded(codes: Vec<u32>, values: &[Option<&str>]) -> Vector {
Vector::dictionary(codes, letters(values)).expect("a dictionary of those strings")
}
fn both_ways(
codes: &[(u32, u32)],
first: &[Option<&str>],
second: &[Option<&str>],
) -> (Chunk, Chunk) {
let coded = Chunk::new(vec![
letters_coded(codes.iter().map(|&(left, _)| left).collect(), first),
letters_coded(codes.iter().map(|&(_, right)| right).collect(), second),
])
.expect("two aligned columns");
let flat = Chunk::new(vec![
letters(&codes.iter().map(|&(left, _)| first[left as usize]).collect::<Vec<_>>()),
letters(&codes.iter().map(|&(_, right)| second[right as usize]).collect::<Vec<_>>()),
])
.expect("two aligned columns");
(coded, flat)
}
#[test]
fn a_group_by_over_two_dictionaries_counts_what_the_flat_columns_count() {
let flags = [Some("A"), Some("N"), Some("R")];
let status = [Some("F"), Some("O")];
let (first_coded, first_flat) =
both_ways(&[(0, 0), (1, 1), (2, 0), (0, 0)], &flags, &status);
let (second_coded, second_flat) = both_ways(&[(1, 1), (1, 0), (2, 0)], &flags, &status);
let counted = two_key_counts(vec![first_coded, second_coded]);
assert_eq!(counted, two_key_counts(vec![first_flat, second_flat]));
assert_eq!(counted.len(), 4, "A/F, N/O, R/F, N/F and nothing else");
}
#[test]
fn a_second_dictionary_does_not_inherit_the_first_one_s_map() {
let first = [Some("A"), Some("N")];
let second = [Some("N"), Some("A")];
let both = [Some("F"), Some("O")];
let (first_coded, first_flat) = both_ways(&[(0, 0), (1, 1)], &first, &both);
let (second_coded, second_flat) = both_ways(&[(0, 0), (1, 1)], &second, &both);
let coded = two_key_counts(vec![first_coded, second_coded]);
assert_eq!(coded, two_key_counts(vec![first_flat, second_flat]));
assert_eq!(coded.len(), 4, "A/F, N/O, N/F and A/O");
}
#[test]
fn null_keys_behind_a_dictionary_group_the_way_flat_nulls_do() {
let flags = [Some("A"), None];
let status = [Some("F"), None];
let (coded, flat) = both_ways(&[(0, 0), (1, 0), (1, 1), (0, 1)], &flags, &status);
let counted = two_key_counts(vec![coded]);
assert_eq!(counted, two_key_counts(vec![flat]));
assert_eq!(counted.len(), 4);
}
fn encoded_count_answer(users: Vector, phrases: Vector) -> Vec<Vec<Value>> {
let plan = Plan::parse(concat!(
"Aggregate #1 groups=[#0.0::BIGINT, #0.1::VARCHAR] ",
"aggregates=[count_star()::BIGINT]\n",
" Get memory.main.t AS t #0 [x::BIGINT, y::VARCHAR]",
))
.expect("a two-key count plan");
let schema = Schema::numbered(
vec![Field::new("x", LogicalType::BigInt), Field::new("y", LogicalType::Varchar)],
0,
);
let rudb_plan::Node::Aggregate { groups, aggregates, .. } = *plan.node(plan.root()) else {
panic!("the root is an aggregate")
};
let (aggregate, out) =
Aggregate::new(&plan, &schema, 1, groups, aggregates, &Memory::unlimited())
.expect("a count aggregate");
let aggregate = aggregate.top_counts(10, 0);
let input = Chunk::new(vec![users, phrases]).expect("two aligned columns");
let mut local = aggregate.local();
aggregate.sink(&input, &mut local).expect("the chunk");
aggregate.combine(local).expect("the one instance");
assert!(aggregate.encoded_count.get().is_some(), "the compact path was selected");
aggregate.finalize(&rudb_pipeline::Lease::alone()).expect("the answer");
let mut rows = answer(&out);
rows.sort_by_key(|row| format!("{row:?}"));
rows
}
#[test]
fn runs_of_one_key_count_every_row_in_them() {
let dictionary = Arc::new(
Vector::from_values(
LogicalType::Varchar,
&[Value::Varchar("one".into()), Value::Varchar("two".into())],
)
.expect("search phrases"),
);
let users = [7, 7, 7, 8, 7, 7, 7].map(Value::BigInt);
let users = Vector::from_values(LogicalType::BigInt, &users).expect("user ids");
let phrases =
Vector::stable_dictionary(vec![0, 0, 1, 1, 0, 0, 0], dictionary).expect("stable codes");
let mut expected = vec![
vec![Value::BigInt(7), Value::Varchar("one".into()), Value::BigInt(5)],
vec![Value::BigInt(7), Value::Varchar("two".into()), Value::BigInt(1)],
vec![Value::BigInt(8), Value::Varchar("two".into()), Value::BigInt(1)],
];
expected.sort_by_key(|row| format!("{row:?}"));
assert_eq!(encoded_count_answer(users, phrases), expected);
}
#[test]
fn a_three_key_count_folds_tucked_and_wide_records_into_the_same_groups() {
let plan = Plan::parse(concat!(
"Aggregate #1 groups=[#0.0::BIGINT, #0.1::BIGINT, #0.2::VARCHAR] ",
"aggregates=[count_star()::BIGINT]\n",
" Get memory.main.t AS t #0 [x::BIGINT, m::BIGINT, y::VARCHAR]",
))
.expect("a three-key count plan");
let schema = Schema::numbered(
vec![
Field::new("x", LogicalType::BigInt),
Field::new("m", LogicalType::BigInt),
Field::new("y", LogicalType::Varchar),
],
0,
);
let rudb_plan::Node::Aggregate { groups, aggregates, .. } = *plan.node(plan.root()) else {
panic!("the root is an aggregate")
};
let (aggregate, out) =
Aggregate::new(&plan, &schema, 1, groups, aggregates, &Memory::unlimited())
.expect("a count aggregate");
let aggregate = aggregate.top_counts(10, 0);
let dictionary = Arc::new(
Vector::from_values(
LogicalType::Varchar,
&[Value::Varchar("one".into()), Value::Varchar("two".into())],
)
.expect("search phrases"),
);
let top = (1_i64 << 31) - 1;
let seconds =
[Value::BigInt(3), Value::BigInt(-1), Value::BigInt(top), Value::BigInt(top + 1)];
let chunk = |null: bool| {
let mut second = seconds.to_vec();
let mut first = vec![Value::BigInt(7); 4];
if null {
second.push(Value::Null);
first.push(Value::BigInt(7));
}
let codes = if null { vec![0, 1, 0, 1, 0] } else { vec![0, 1, 0, 1] };
Chunk::new(vec![
Vector::from_values(LogicalType::BigInt, &first).expect("user ids"),
Vector::from_values(LogicalType::BigInt, &second).expect("minutes"),
Vector::stable_dictionary(codes, Arc::clone(&dictionary)).expect("stable codes"),
])
.expect("three aligned columns")
};
let mut local = aggregate.local();
aggregate.sink(&chunk(true), &mut local).expect("the chunk with a null");
aggregate.sink(&chunk(false), &mut local).expect("the chunk without one");
aggregate.combine(local).expect("the one instance");
assert!(aggregate.encoded_count.get().is_some(), "the compact path was selected");
aggregate.finalize(&rudb_pipeline::Lease::alone()).expect("the answer");
let mut rows = answer(&out);
rows.sort_by_key(|row| format!("{row:?}"));
let group = |second: Value, phrase: &str, count| {
vec![Value::BigInt(7), second, Value::Varchar(phrase.into()), Value::BigInt(count)]
};
let mut expected = vec![
group(Value::BigInt(3), "one", 2),
group(Value::BigInt(-1), "two", 2),
group(Value::BigInt(top), "one", 2),
group(Value::BigInt(top + 1), "two", 2),
group(Value::Null, "one", 1),
];
expected.sort_by_key(|row| format!("{row:?}"));
assert_eq!(rows, expected);
}
#[test]
fn a_tucked_second_key_comes_back_apart_from_its_code() {
assert_eq!(Tucked.keys(5 << 21 | 12_345, 21), (5, 12_345));
assert_eq!(Tucked.keys(u32::MAX, 32), (0, u32::MAX));
assert_eq!(Tucked.keys(9, 0), (9, 0));
assert_eq!(super::code_bits(2), 1);
assert_eq!(super::code_bits(1), 0);
assert_eq!(super::tuck_limit(32), 0);
assert_eq!(super::tuck_limit(0), 1 << 32);
}
#[test]
fn the_run_reader_and_the_row_at_a_time_scatter_count_the_same_groups() {
let phrases = || {
let dictionary = Arc::new(
Vector::from_values(
LogicalType::Varchar,
&[Value::Varchar("one".into()), Value::Varchar("two".into())],
)
.expect("search phrases"),
);
Vector::stable_dictionary(vec![0, 0, 1], dictionary).expect("stable codes")
};
let values = [Value::BigInt(7), Value::BigInt(7), Value::BigInt(8)];
let flat = Vector::from_values(LogicalType::BigInt, &values).expect("user ids");
let coded = Vector::stable_dictionary(
vec![0, 0, 1],
Arc::new(
Vector::from_values(LogicalType::BigInt, &[Value::BigInt(7), Value::BigInt(8)])
.expect("distinct user ids"),
),
)
.expect("coded user ids");
assert!(Signed::of(&flat, 3).is_some(), "a flat key is read as a run of words");
assert!(Signed::of(&coded, 3).is_none(), "a coded key goes down the row at a time loop");
assert_eq!(encoded_count_answer(flat, phrases()), encoded_count_answer(coded, phrases()));
}
#[test]
fn a_null_leading_key_still_gets_a_group_of_its_own() {
let dictionary = Arc::new(
Vector::from_values(LogicalType::Varchar, &[Value::Varchar("one".into())])
.expect("search phrases"),
);
let phrases = Vector::stable_dictionary(vec![0, 0, 0], dictionary).expect("stable codes");
let users =
Vector::from_values(LogicalType::BigInt, &[Value::BigInt(7), Value::Null, Value::Null])
.expect("user ids");
assert!(Signed::of(&users, 3).is_none(), "a key with a null is not read as a run of words");
assert_eq!(
encoded_count_answer(users, phrases),
vec![
vec![Value::BigInt(7), Value::Varchar("one".into()), Value::BigInt(1)],
vec![Value::Null, Value::Varchar("one".into()), Value::BigInt(2)],
]
);
}
#[test]
fn fixed_radix_partition_aggregates_collisions_and_nulls_exactly() {
let mut partition = FixedRun::default();
let row = |first, second, sum, mean| FixedRecord { first, second, sum, mean };
partition.push(row(1, 2, 3, 4), FixedRecord::ALL);
partition
.push(row(1, 2, 5, 0), FixedRecord::FIRST | FixedRecord::SECOND | FixedRecord::SUM);
partition.push(row(0, 2, 0, 6), FixedRecord::SECOND | FixedRecord::MEAN);
partition.push(row(0, 2, 7, 8), FixedRecord::SECOND | FixedRecord::SUM | FixedRecord::MEAN);
let calls = [
Call {
name: "count_star".into(),
args: Vec::new(),
distinct: false,
filter: None,
returns: LogicalType::BigInt,
affine: None,
reads_total: None,
repeats: None,
},
Call {
name: "sum".into(),
args: Vec::new(),
distinct: false,
filter: None,
returns: LogicalType::HugeInt,
affine: None,
reads_total: None,
repeats: None,
},
Call {
name: "avg".into(),
args: Vec::new(),
distinct: false,
filter: None,
returns: LogicalType::Double,
affine: None,
reads_total: None,
repeats: None,
},
];
let keys = [LogicalType::SmallInt, LogicalType::Integer];
let part = fixed_partition(
&mut FixedRuns { runs: vec![partition] },
&keys,
10,
&calls,
&Memory::unlimited(),
)
.expect("the fixed partition");
let mut rows: Vec<Vec<Value>> = Vec::new();
for chunk in part.chunks {
for row in 0..chunk.len() {
rows.push((0..chunk.width()).map(|column| chunk.value_at(row, column)).collect());
}
}
rows.sort_by_key(|row| format!("{:?}", row[0]));
assert_eq!(
rows,
[
vec![
Value::Null,
Value::Integer(2),
Value::BigInt(2),
Value::HugeInt(7),
Value::Double(7.0),
],
vec![
Value::SmallInt(1),
Value::Integer(2),
Value::BigInt(2),
Value::HugeInt(8),
Value::Double(4.0),
],
]
);
assert_eq!(size_of::<FixedRecord>(), 16);
}
#[test]
fn a_fixed_partition_split_for_cache_keeps_the_largest_groups_of_every_split() {
let groups = super::FIXED_SPLIT_ROWS as i64 * 8;
let row = |first: i64| FixedRecord { first, second: (first % 7) as i32, sum: 1, mean: 2 };
let (mut early, mut late) = (FixedRun::default(), FixedRun::default());
for first in 0..groups {
let times = if first % 5_000 == 0 { 3 + first / 5_000 } else { 1 };
for time in 0..times {
let into = if time % 2 == 0 { &mut early } else { &mut late };
into.push(row(first), FixedRecord::ALL);
}
}
let call = |name: &str, returns| Call {
name: name.into(),
args: Vec::new(),
distinct: false,
filter: None,
returns,
affine: None,
reads_total: None,
repeats: None,
};
let calls = [
call("count_star", LogicalType::BigInt),
call("sum", LogicalType::HugeInt),
call("avg", LogicalType::Double),
];
let part = fixed_partition(
&mut FixedRuns { runs: vec![early, late] },
&[LogicalType::BigInt, LogicalType::Integer],
4,
&calls,
&Memory::unlimited(),
)
.expect("the fixed partition");
let mut rows: Vec<Vec<Value>> = Vec::new();
for chunk in part.chunks {
for row in 0..chunk.len() {
rows.push((0..chunk.width()).map(|column| chunk.value_at(row, column)).collect());
}
}
let largest = groups / 5_000 * 5_000;
let expected = (0..4)
.map(|step| {
let first = largest - step * 5_000;
let times = 3 + first / 5_000;
vec![
Value::BigInt(first),
Value::Integer((first % 7) as i32),
Value::BigInt(times),
Value::HugeInt(i128::from(times)),
Value::Double(2.0),
]
})
.collect::<Vec<_>>();
assert_eq!(rows, expected, "the four largest groups, largest first");
}
#[test]
fn the_largest_counts_come_back_largest_first_and_in_slot_order_among_equals() {
let counts = [3, 1, 5, 3, 5, 1, 2, 3];
assert_eq!(super::largest(counts.len(), 4, |slot| counts[slot]), [2, 4, 0, 3]);
assert_eq!(super::largest(counts.len(), 0, |slot| counts[slot]), [0_usize; 0]);
assert_eq!(super::largest(counts.len(), 20, |slot| counts[slot]), [2, 4, 0, 3, 7, 6, 1, 5]);
assert_eq!(super::largest(counts.len(), 1, |slot| counts[slot]), [2]);
}
#[test]
fn a_dense_partition_sorts_a_few_rows_into_the_answer_the_array_gives() {
let spellings =
(0..1_000).map(|code| Value::Varchar(format!("v{code}"))).collect::<Vec<_>>();
let dictionary =
Arc::new(Vector::from_values(LogicalType::Varchar, &spellings).expect("a dictionary"));
let bounded = |runs: Vec<Vec<u32>>, bound: Option<usize>| {
let runs = runs.into_iter().map(|run| run.into_iter().collect()).collect();
let mut partition = super::DensePartition { runs, nulls: 0 };
let chunks = super::dense_partition(
&dictionary,
1,
&mut partition,
bound,
&[None],
&[LogicalType::Varchar],
)
.expect("the dense partition");
let mut rows: Vec<(Value, Value)> = Vec::new();
for chunk in chunks {
for row in 0..chunk.len() {
rows.push((chunk.value_at(row, 0), chunk.value_at(row, 1)));
}
}
rows
};
let answer = |runs: Vec<Vec<u32>>| bounded(runs, None);
let codes = [405, 9, 5, 9, 405, 405, 997];
let few = answer(vec![codes[..4].to_vec(), codes[4..].to_vec()]);
let many = answer(vec![codes.repeat(100)]);
let expected = |times: i64| {
[(5, 1), (9, 2), (405, 3), (997, 1)]
.map(|(code, count)| {
(Value::Varchar(format!("v{code}")), Value::BigInt(count * times))
})
.to_vec()
};
assert_eq!(few, expected(1));
assert_eq!(many, expected(100));
let top = |times: i64| {
[(9, 2), (405, 3)]
.map(|(code, count)| {
(Value::Varchar(format!("v{code}")), Value::BigInt(count * times))
})
.to_vec()
};
assert_eq!(bounded(vec![codes.to_vec()], Some(2)), top(1));
assert_eq!(bounded(vec![codes.repeat(100)], Some(2)), top(100));
assert_eq!(bounded(vec![codes.repeat(100)], Some(4)), expected(100));
}
#[test]
fn the_one_table_an_aggregate_answers_out_of_takes_room_for_every_group() {
assert_eq!(Share::Whole.of(8 << 20), Some(8 << 20));
assert_eq!(Share::Whole.of(1), Some(1));
}
#[test]
fn a_partition_takes_room_for_its_share() {
let groups = 8 << 20;
assert_eq!(Share::Partition.of(groups), Some(groups / RADIX_PARTITIONS as u64));
assert_eq!(Share::Partition.of(6400), Some(100));
}
#[test]
fn a_partition_of_a_small_aggregate_still_gets_a_group() {
for groups in 0..RADIX_PARTITIONS as u64 {
assert_eq!(Share::Partition.of(groups), Some(1), "{groups} groups");
}
}
#[test]
fn the_tables_that_pass_the_groups_on_take_no_more_room_than_they_will_use() {
assert_eq!(Share::Passing.of(8 << 20), Some(PARTITION_FROM as u64));
assert_eq!(Share::Passing.of(100), Some(100));
assert_eq!(Share::Local.of(8 << 20), None);
assert_eq!(Share::Local.of(1), None);
}
fn folded(aggregate: &Aggregate<'_>, chunks: &[Vec<i32>]) -> super::Building {
let mut local = aggregate.local();
let mut building = aggregate.starting(Share::Local);
for part in chunks {
let rows = aggregate.read(&chunk(part), &mut local.expressions).expect("a chunk");
aggregate.fold(&rows, &mut building, None, None).expect("folded");
}
building
}
#[test]
fn a_key_that_keeps_moving_past_its_window_builds_maps_only_as_the_rows_pay_for_them() {
let plan = parsed("Aggregate #1 groups=[#0.0::INTEGER] aggregates=[count_star()::BIGINT]");
let (aggregate, _out) = aggregate(&plan);
let chunks: Vec<Vec<i32>> =
(0..100).map(|at| (at * 2_048..(at + 1) * 2_048).collect()).collect();
let building = folded(&aggregate, &chunks);
assert_eq!(building.coded_read, 100 * 2_048);
assert!(building.coded_spent > 0, "the first chunks are read by value");
assert!(
building.coded_spent <= building.coded_read * WINDOW_RATE + WINDOW_SLACK,
"{} places cleared for {} rows",
building.coded_spent,
building.coded_read
);
}
#[test]
fn a_key_that_stays_inside_its_window_keeps_its_map() {
let plan = parsed("Aggregate #1 groups=[#0.0::INTEGER] aggregates=[count_star()::BIGINT]");
let (aggregate, _out) = aggregate(&plan);
let chunks: Vec<Vec<i32>> = (0..100)
.map(|at| (0..2_048).map(|row| (row * 7 + at * 13) % 3_000 + 17).collect())
.collect();
let building = folded(&aggregate, &chunks);
assert!(!building.coded_on.is_empty(), "the last chunk was answered by the map");
assert!(building.coded_spent < 64 * 1_024, "the window settled after a few builds");
}
#[test]
fn a_key_that_climbs_grows_its_map_in_place() {
let plan = parsed("Aggregate #1 groups=[#0.0::INTEGER] aggregates=[count_star()::BIGINT]");
let (aggregate, _out) = aggregate(&plan);
let chunks: Vec<Vec<i32>> =
(0..32).map(|at| (at * 1_024..(at + 1) * 1_024).collect()).collect();
let building = folded(&aggregate, &chunks);
assert!(!building.coded_on.is_empty(), "the last chunk was answered by the map");
assert!(building.coded_map.len() > 32 * 1_024, "the map covers every value");
assert_eq!(building.coded_spent, building.coded_map.len(), "no place was cleared twice");
}
#[test]
fn a_direct_index_over_the_range_goes_on_the_table_that_sees_the_whole_range() {
assert!(Share::Whole.before_the_split());
assert!(Share::Passing.before_the_split());
assert!(!Share::Partition.before_the_split());
assert!(!Share::Local.before_the_split());
}
#[test]
fn a_grouped_key_closes_its_inner_runs_whichever_way_they_go() {
let values = [9, 9, 4, 4, 4, 7, 2, 2, 5].map(Value::BigInt);
let key = Vector::from_values(LogicalType::BigInt, &values).expect("keys");
assert_eq!(interior(&key, values.len(), false), None, "ascending refuses it");
assert_eq!(
interior(&key, values.len(), true),
Some((2, 8)),
"grouped closes 4, 7 and 2, and leaves 9 and 5 to the table"
);
let two = Vector::from_values(LogicalType::BigInt, &values[..5]).expect("keys");
assert_eq!(interior(&two, 5, true), None, "two runs have nothing strictly inside");
}
}