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, Reservation, Result, Session, Stage, Value, stage,
};
use rudb_kernels::{Accumulator, NOWHERE, is_true, update_scattered};
use rudb_pipeline::{Lease, Progress, Sink};
use rudb_plan::{Expr, ExprRef, Plan, Slice};
use rudb_vector::{Chunk, Data, Form, VECTOR_SIZE, Validity, Vector};
use crate::buffer::Buffered;
use crate::group_distinct;
use crate::group_mixed;
use crate::key::{BigIntSet, Key, RowSet, mix, spread};
use crate::pairs::together;
use crate::prepared::{Prepared, Scratch};
use crate::rows;
use crate::schema::Schema;
use crate::spill::{Reader, Spill};
use crate::table::{Probe, Table, Walk};
#[derive(Debug, Clone)]
struct Call {
name: String,
args: Vec<ExprRef>,
distinct: bool,
filter: Option<ExprRef>,
returns: LogicalType,
affine: Option<(usize, i64)>,
}
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));
}
}
}
#[derive(Debug)]
pub(crate) struct Aggregate<'a> {
plan: &'a Plan,
keys: Vec<ExprRef>,
groups: Vec<ExprRef>,
constants: Vec<Option<Value>>,
calls: Vec<Call>,
inputs: Prepared,
schema: Schema,
alone: bool,
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>,
having_count: Option<(usize, i64)>,
max_groups: Option<usize>,
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<EncodedCountExchange>>,
grouped_distinct: OnceLock<Option<group_distinct::Exchange>>,
mixed: OnceLock<group_mixed::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<Vec<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 {
dictionary: Arc<Vector>,
leading: Vec<LogicalType>,
dictionary_nulls: bool,
partitions: Vec<Mutex<EncodedCountRuns>>,
held: Mutex<Vec<Reservation>>,
}
#[derive(Debug, Default)]
struct EncodedCountRuns {
runs: Vec<EncodedCountPartition>,
}
#[derive(Debug, Default)]
struct FixedRuns {
runs: Vec<FixedPartition>,
}
impl EncodedCountRuns {
fn seed(&mut self) -> (EncodedCountPartition, usize) {
let total = self.runs.iter().map(|run| run.rows.len()).sum();
let widest = widest_run(self.runs.iter().map(|run| run.rows.len()));
let seed = widest.map(|at| self.runs.swap_remove(at)).unwrap_or_default();
(seed, total)
}
}
impl FixedRuns {
fn seed(&mut self) -> (FixedPartition, usize) {
let total = self.runs.iter().map(|run| run.rows.len()).sum();
let widest = widest_run(self.runs.iter().map(|run| run.rows.len()));
let seed = widest.map(|at| self.runs.swap_remove(at)).unwrap_or_default();
(seed, total)
}
}
fn widest_run(lengths: impl Iterator<Item = usize>) -> Option<usize> {
lengths.enumerate().max_by_key(|&(_, rows)| rows).map(|(at, _)| at)
}
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 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 {
first: i64,
second: i64,
hash: u32,
third: u32,
}
#[derive(Debug, Default)]
struct EncodedCountPartition {
rows: Vec<EncodedCountRecord>,
validity: Vec<u8>,
}
impl EncodedCountRecord {
const FIRST: u8 = 1;
const SECOND: u8 = 2;
const THIRD: u8 = 4;
const ALL: u8 = Self::FIRST | Self::SECOND | Self::THIRD;
}
impl EncodedCountPartition {
fn push(&mut self, row: EncodedCountRecord, valid: u8) {
let keeping = !self.validity.is_empty() || valid != EncodedCountRecord::ALL;
if keeping {
self.validity.resize(self.rows.len(), EncodedCountRecord::ALL);
}
self.rows.push(row);
if keeping {
self.validity.push(valid);
}
}
fn footprint(&self) -> usize {
self.rows.capacity() * size_of::<EncodedCountRecord>()
+ self.validity.capacity() * size_of::<u8>()
}
}
#[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);
}
}
fn footprint(&self) -> usize {
self.rows.capacity() * size_of::<FixedRecord>() + 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 = 16;
const DENSE_PARTITIONS: usize = 4;
const PARTITION_FROM: usize = 4_096;
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,
});
}
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 mut inputs = keys.clone();
for call in &calls {
if call.affine.is_none() {
inputs.extend_from_slice(&call.args);
}
inputs.extend(call.filter);
}
let inputs = Prepared::shared(plan, &inputs, &input_schema)?;
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();
let by_vector: Vec<bool> =
calls.iter().map(|call| alone && !call.distinct && call.filter.is_none()).collect();
let out = Buffered::new();
let aggregate = Self {
plan,
keys,
constants,
alone,
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,
agreed: Mutex::new(None),
settled: AtomicBool::new(false),
by_vector,
groups,
calls,
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(),
out: out.clone(),
};
Ok((aggregate, out))
}
pub(crate) fn limit_groups(mut self, max_groups: usize) -> Self {
self.max_groups = Some(max_groups);
self
}
#[must_use]
pub(crate) fn top_counts(mut self, bound: usize) -> Self {
if self.count_only
|| self.compact_numeric
|| self.distinct_count
|| self.mixed_numeric_distinct
{
self.top_counts = Some(bound);
}
self
}
#[must_use]
pub(crate) fn having_count(mut self, call: usize, minimum: i64) -> Self {
if self.calls.get(call).is_some_and(|call| {
call.name == "count_star"
&& call.args.is_empty()
&& !call.distinct
&& call.filter.is_none()
}) {
self.having_count = Some((call, minimum));
}
self
}
pub(crate) fn schema(&self) -> &Schema {
&self.schema
}
fn read(&self, chunk: &Chunk, scratch: &mut Scratch) -> Result<Rows> {
let rows = chunk.len();
let mut evaluated = Vec::new();
self.inputs.evaluate(chunk, scratch, &mut evaluated)?;
let mut values = evaluated.into_iter();
let keys = values.by_ref().take(self.keys.len()).collect();
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.affine.is_none() {
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 })
}
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)
&& self.keys.len() == 1
&& matches!(
self.plan.expr_type(self.keys[0]),
LogicalType::Integer | LogicalType::Varchar
)
}
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 [FixedPartition],
memory: &mut Reservation,
) -> 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(FixedPartition::footprint).sum::<usize>();
let shift = u64::BITS - RADIX_PARTITIONS.ilog2();
for row in 0..rows.rows {
let mut valid = 0;
let first_value = if first.is_null_at(row) {
0
} else {
valid |= FixedRecord::FIRST;
i64::try_from(first.signed_at(row).ok_or_else(|| {
Error::internal("a fixed first key has no signed representation")
})?)
.map_err(|_| Error::internal("a fixed first key is out of range"))?
};
let second_value = if second.is_null_at(row) {
0
} else {
valid |= FixedRecord::SECOND;
i32::try_from(second.signed_at(row).ok_or_else(|| {
Error::internal("a fixed second key has no signed representation")
})?)
.map_err(|_| Error::internal("a fixed second key is out of range"))?
};
let sum_value = if sum.is_null_at(row) {
0
} else {
valid |= FixedRecord::SUM;
i16::try_from(sum.signed_at(row).ok_or_else(|| {
Error::internal("a fixed SMALLINT sum has no signed representation")
})?)
.map_err(|_| Error::internal("a fixed SMALLINT sum is out of range"))?
};
let mean_value = if mean.is_null_at(row) {
0
} else {
valid |= FixedRecord::MEAN;
i16::try_from(mean.signed_at(row).ok_or_else(|| {
Error::internal("a fixed SMALLINT mean has no signed representation")
})?)
.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(FixedPartition::footprint).sum::<usize>();
memory.grow(width_of(after.saturating_sub(before)))
}
fn buffer_encoded_count(
&self,
rows: &Rows,
partitions: &mut [EncodedCountPartition],
memory: &mut Reservation,
) -> Result<bool> {
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 = third.stable_dictionary_parts();
let state = self.encoded_count.get_or_init(|| {
dictionary.as_ref().map(|(_, dictionary)| EncodedCountExchange {
dictionary: Arc::clone(dictionary),
leading: self.keys[..self.keys.len() - 1]
.iter()
.map(|&key| self.plan.expr_type(key).clone())
.collect(),
dictionary_nulls: dictionary.validity().has_nulls(dictionary.len())
|| !nulls_are_in_the_mask(dictionary),
partitions: (0..RADIX_PARTITIONS)
.map(|_| Mutex::new(EncodedCountRuns::default()))
.collect(),
held: Mutex::new(Vec::new()),
})
});
let Some(state) = state else { return Ok(false) };
let Some((codes, dictionary)) = dictionary else {
return Err(Error::internal(
"an encoded count exchange changed from dictionary to flat strings",
));
};
if !Arc::ptr_eq(&state.dictionary, dictionary) {
return Err(Error::internal(
"an encoded count exchange received two string code spaces",
));
}
if state.leading.len() + 1 != rows.keys.len() {
return Err(Error::internal("an encoded count exchange changed key width"));
}
let before = partitions.iter().map(EncodedCountPartition::footprint).sum::<usize>();
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(|_| {
!state.dictionary_nulls
&& third.len() >= rows.rows
&& !third.validity().has_nulls(rows.rows)
});
if let Some((first, second)) = plain {
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));
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;
partitions[(hash >> shift) as usize].push(
EncodedCountRecord {
first: first_value,
second: second_value,
hash,
third: third_code,
},
EncodedCountRecord::ALL,
);
}
let after = partitions.iter().map(EncodedCountPartition::footprint).sum::<usize>();
memory.grow(width_of(after.saturating_sub(before)))?;
return Ok(true);
}
for (row, &third_code) in codes.iter().enumerate().take(rows.rows) {
let mut valid = if second.is_none() { EncodedCountRecord::SECOND } else { 0 };
let first_value = if first.is_null_at(row) {
0
} else {
valid |= EncodedCountRecord::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 |= EncodedCountRecord::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 |= EncodedCountRecord::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 & EncodedCountRecord::FIRST != 0 { first_value as u64 } else { NOTHING };
let second_word =
if valid & EncodedCountRecord::SECOND != 0 { second_value as u64 } else { NOTHING };
let third_word = if valid & EncodedCountRecord::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[(hash >> shift) as usize].push(
EncodedCountRecord {
first: first_value,
second: second_value,
hash,
third: third_value,
},
valid,
);
}
let after = partitions.iter().map(EncodedCountPartition::footprint).sum::<usize>();
memory.grow(width_of(after.saturating_sub(before)))?;
Ok(true)
}
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.start();
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);
timing.stop(0);
folded?;
}
self.finish(local, chunks, held)
}
fn start(&self) -> Building {
let calls = self.calls.len();
let mut local = Building {
scratch: self.memory.reservation(),
containers: self.memory.reservation(),
charged: 0,
charged_keys: 0,
table: Table::new(
&self.keys.iter().map(|&key| self.plan.expr_type(key).clone()).collect::<Vec<_>>(),
),
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(),
kept: Vec::new(),
affine_rows: vec![0; calls],
over: None,
away: Vec::new(),
failure: None,
};
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]>,
) -> Result<()> {
let Building {
scratch,
containers,
charged,
charged_keys,
table,
states,
counts,
compact,
overflow,
seen,
groups,
given,
hashes,
slots,
walk,
kept,
affine_rows,
over,
away,
failure: _,
} = local;
let calls = self.calls.len();
let alone = self.alone;
let Rows { keys, arguments, filters, rows: length } = seen_rows;
let mut aside = 0;
for at in 0..calls {
if self.calls[at].affine.is_some() {
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(());
}
if !alone {
match prehashed {
Some(prehashed) => {
hashes.clear();
hashes.extend_from_slice(prehashed);
}
None => crate::table::hash(keys, *length, hashes, crate::table::Across::OneInput),
}
}
slots.clear();
slots.resize(*length, if alone { 0 } else { NOWHERE });
let mut from = 0;
while !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);
}
}
}
if self.compact_numeric {
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,
)?;
}
}
if self.count_only {
for &slot in slots.iter() {
if slot != NOWHERE {
counts[slot] += 1;
}
}
}
for (at, call) in self.calls.iter().enumerate() {
if self.count_only || self.compact_numeric {
break;
}
if self.by_vector[at] || call.affine.is_some() {
continue;
}
if call.distinct {
aside += self.distinct(states, seen, seen_rows, slots, at, given)?;
continue;
}
let picked = match &filters[at] {
None => &*slots,
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
}
};
update_scattered(states, picked, calls, at, arguments[at].first(), *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 && 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,
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), _) => {
let mut best = Vec::with_capacity(bound.min(groups));
for slot in 0..groups {
let at = best.partition_point(|&kept| count(kept, 0) >= count(slot, 0));
if at < bound {
best.insert(at, slot);
best.truncate(bound);
}
}
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);
scratch.grow(width_of(VECTOR_SIZE.min(output_groups) * size_of::<Value>()))?;
let mut results: Vec<Value> = 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 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 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 {
Some((source, offset)) => states[slot * calls + source]
.finish_offset(offset, affine_rows[source]),
None => states[slot * calls + at].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;
for &slot in slots {
let row = slot - start;
let hash = source.hash_of(slot);
let target = match into.table.probe(hash, keys, row) {
Probe::Found(target) => target,
Probe::Vacant(bucket) => {
if self.max_groups.is_some_and(|limit| into.table.len() >= limit) {
continue;
}
let target = into.table.insert(bucket, hash, 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);
}
target
}
};
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.start());
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.start());
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 partition_from(&self) -> usize {
match self.top_counts {
Some(bound) => PARTITION_FROM.max(bound.saturating_mul(self.merged.len())),
None => PARTITION_FROM,
}
}
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, spreading: &mut Spreading, own: &mut [Option<Building>]) -> Result<bool> {
let spilled = own.iter().flatten().any(|table| table.over.is_some());
if !spilled && !crowded(&self.memory) && self.room_for_local(own) {
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, own: &[Option<Building>]) -> bool {
let Some(limit) = self.memory.limit() else { return true };
let mine: u64 = own
.iter()
.flatten()
.map(|table| table.scratch.bytes() + table.containers.bytes())
.sum();
let instances = self.started.load(Ordering::Relaxed) as u64;
mine.saturating_mul(instances) < limit / 4
}
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 split(&self, rows: &Rows, spreading: &mut Spreading) -> Result<Vec<Option<Rows>>> {
let Spreading { hashes, picks, keyed, spin, .. } = 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();
for (row, &hash) in hashes.iter().enumerate() {
let partition = (hash >> shift) as usize;
picks[partition].push(row as u32);
keyed[partition].push(hash);
}
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.start());
if let Some(error) = table.failure.take() {
return Err(error);
}
self.fold(selected, table, Some(&spreading.keyed[partition]))?;
}
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.start());
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.start());
self.fold(selected, table, Some(&keyed[partition]))?;
}
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.start());
self.fold(selected, table, Some(&keyed[partition]))?;
}
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] {
if !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(());
}
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 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
}
}
struct Rows {
keys: Vec<Vector>,
arguments: Vec<Vec<Vector>>,
filters: Vec<Option<Vector>>,
rows: usize,
}
impl Rows {
fn width(&self) -> usize {
self.keys.len()
+ self.arguments.iter().map(Vec::len).sum::<usize>()
+ self.filters.iter().flatten().count()
}
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(),
})
}
}
#[derive(Debug)]
struct Agreed {
table: Table,
hashes: Vec<u64>,
keys: Option<Arc<Vec<Vector>>>,
}
#[derive(Debug)]
pub(crate) struct Partitioned {
mixed: group_mixed::Local,
grouped_distinct: group_distinct::Local,
encoded: bool,
encoded_records: Vec<EncodedCountPartition>,
encoded_memory: Reservation,
radix_distinct: bool,
radix_distinct_records: Vec<BigIntDistinctPartition>,
radix_distinct_memory: Reservation,
fixed: bool,
fixed_records: Vec<FixedPartition>,
fixed_memory: Reservation,
dense: bool,
dense_codes: Vec<Vec<u32>>,
dense_nulls: i64,
dense_memory: Reservation,
single: Option<Building>,
installed: bool,
expressions: Scratch,
spreading: Spreading,
own: Vec<Option<Building>>,
}
#[derive(Debug)]
struct Spreading {
hashes: Vec<u64>,
picks: Vec<Vec<u32>>,
keyed: Vec<Vec<u64>>,
spin: usize,
waiting: Vec<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(),
}
}
}
#[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,
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) {
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 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 }))
}
}
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 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 local(&self) -> Partitioned {
self.started.fetch_add(1, Ordering::Relaxed);
Partitioned {
mixed: group_mixed::Local::new(&self.memory),
grouped_distinct: group_distinct::Local::new(&self.memory),
encoded: false,
encoded_records: (0..RADIX_PARTITIONS)
.map(|_| EncodedCountPartition::default())
.collect(),
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(|_| FixedPartition::default()).collect(),
fixed_memory: self.memory.reservation(),
dense: false,
dense_codes: vec![Vec::new(); DENSE_PARTITIONS],
dense_nulls: 0,
dense_memory: self.memory.reservation(),
single: Some(self.start()),
installed: false,
expressions: self.inputs.scratch(),
spreading: Spreading::new(),
own: (0..RADIX_PARTITIONS).map(|_| None).collect(),
}
}
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,
grouped_distinct,
encoded,
encoded_records,
encoded_memory,
radix_distinct,
radix_distinct_records,
radix_distinct_memory,
fixed,
fixed_records,
fixed_memory,
dense,
dense_codes,
dense_nulls,
dense_memory,
single,
installed,
expressions,
spreading,
own,
} = local;
let rows = self.read(chunk, expressions)?;
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.grouped_distinct_top_count() {
let [group] = rows.keys.as_slice() else {
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 codes = if self.plan.expr_type(self.keys[0]) == &LogicalType::Varchar {
match group.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
};
let timing = stage::Timing::start(Stage::Scatter);
let buffered = group_distinct::Exchange::buffer(
&self.grouped_distinct,
group,
codes,
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);
timing.stop(0);
buffered?;
*fixed = true;
return Ok(Progress::More);
}
if self.count_only && self.keys.len() == 1 {
if let [key] = rows.keys.as_slice() {
if 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(Vec::capacity).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(Vec::capacity).sum::<usize>();
dense_memory.grow(width_of(after.saturating_sub(before) * size_of::<u32>()))?;
*dense = true;
return Ok(Progress::More);
}
}
}
if let Some(table) = single {
if let Some(error) = table.failure.take() {
return Err(error);
}
if let Some(limit) = self.max_groups {
if !self.alone {
self.agree(&rows, limit, table, installed)?;
}
}
let timing = stage::Timing::start(Stage::Fold);
let folded = self.fold(&rows, table, None);
timing.stop(0);
folded?;
if !self.ought_to_partition(table) {
return Ok(Progress::More);
}
let handing = single.take().expect("the table was there a moment ago");
self.begin_partitioning(spreading, own)?;
self.hand(handing, spreading, own)?;
return Ok(Progress::More);
}
if self.locally.load(Ordering::Relaxed) && self.still_local(spreading, own)? {
self.spread_own(&rows, spreading, own)?;
return Ok(Progress::More);
}
self.spread(&rows, spreading)?;
Ok(Progress::More)
}
fn combine(&self, local: Partitioned) -> Result<()> {
let Partitioned {
mixed,
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,
mut spreading,
mut own,
..
} = local;
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 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");
for (partition, rows) in encoded_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(encoded_memory);
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);
}
let mut kept = self.merged[0].lock().map_err(poisoned)?;
if kept.table.is_none() {
kept.table = Some(arriving);
return Ok(());
}
let spilled =
arriving.over.is_some() || kept.table.as_ref().is_some_and(|held| held.over.is_some());
if spilled {
let seeded = kept.table.take().expect("the table was there a moment ago");
built.partitioning = true;
drop(kept);
drop(built);
self.hand_over(seeded, &mut spreading)?;
return self.hand_over(arriving, &mut spreading);
}
self.merge(arriving, kept.table.as_mut().expect("the table was there a moment ago"))?;
Ok(())
}
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"),
&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"),
&self.memory,
)?;
return self.out.fill(chunks);
}
if let Some(Some(encoded)) = self.encoded_count.get() {
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(|run| run.rows.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");
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);
return 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");
let next = AtomicUsize::new(0);
let slots: Vec<Mutex<Option<Result<Part>>>> =
(0..RADIX_PARTITIONS).map(|_| Mutex::new(None)).collect();
let degree = threads.degree().clamp(1, RADIX_PARTITIONS);
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.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(
next: &AtomicUsize,
slots: &[Mutex<Option<Result<Part>>>],
encoded: &EncodedCountExchange,
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, bound, memory)
});
if let Ok(mut slot) = slots[at].lock() {
*slot = Some(done);
}
}
}
#[inline]
fn encoded_slot(
buckets: &[u32],
groups: &EncodedCountPartition,
row: EncodedCountRecord,
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 { EncodedCountRecord::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;
}
}
fn encoded_count_partition(
runs: &mut EncodedCountRuns,
dictionary: &Vector,
leading: &[LogicalType],
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 (mut partition, total) = runs.seed();
let capacity = total.saturating_mul(2).max(64).next_power_of_two();
let mut working = memory.reservation();
let room = total.saturating_sub(partition.rows.len());
working.grow(width_of(
capacity * size_of::<u32>()
+ total * size_of::<i64>()
+ room * size_of::<EncodedCountRecord>(),
))?;
let mut buckets = vec![EMPTY_SLOT; capacity];
let mut counts: Vec<i64> = Vec::with_capacity(total);
partition.rows.reserve(room);
reserving.stop(0);
let timing = stage::Timing::start(Stage::Fold);
let seeded = partition.rows.len();
let all_valid = partition.validity.is_empty();
for source in 0..seeded {
let row = partition.rows[source];
let valid = if all_valid { EncodedCountRecord::ALL } else { partition.validity[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(1)
.ok_or_else(|| Error::out_of_range("a grouped COUNT overflowed BIGINT"))?;
}
partition.rows.truncate(counts.len());
if !all_valid {
partition.validity.truncate(counts.len());
}
timing.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();
for (source, &row) in run.rows.iter().enumerate() {
let valid = if all_valid { EncodedCountRecord::ALL } else { run.validity[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.push(row, valid);
counts.push(0);
slot
}
};
counts[slot] = counts[slot]
.checked_add(1)
.ok_or_else(|| Error::out_of_range("a grouped COUNT overflowed BIGINT"))?;
}
}
let all_valid = partition.validity.is_empty();
timing.stop(0);
let timing = stage::Timing::start(Stage::Emit);
let mut best = Vec::with_capacity(bound.min(counts.len()));
for slot in 0..counts.len() {
let at = best.partition_point(|&kept| counts[kept] >= counts[slot]);
if at < bound {
best.insert(at, slot);
best.truncate(bound);
}
}
best.sort_unstable();
let mut output = Vec::with_capacity(best.len());
for slot in best {
let key = partition.rows[slot];
let valid = if all_valid { EncodedCountRecord::ALL } else { partition.validity[slot] };
let first = if valid & EncodedCountRecord::FIRST != 0 {
signed_value(&leading[0], key.first)?
} else {
Value::Null
};
let third = if valid & EncodedCountRecord::THIRD != 0 {
dictionary.try_value_at(key.third as usize)?
} else {
Value::Null
};
let mut row = Vec::with_capacity(keys + 1);
row.push(first);
if keys == 3 {
row.push(if valid & EncodedCountRecord::SECOND != 0 {
signed_value(&leading[1], key.second)?
} else {
Value::Null
});
}
row.push(third);
row.push(Value::BigInt(counts[slot]));
output.push(row);
}
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 bigint_distinct_partition(partition: &mut BigIntDistinctRuns, memory: &Memory) -> Result<i64> {
let held: usize = partition.runs.iter().map(Vec::len).sum();
let capacity = held.saturating_mul(2).max(64).next_power_of_two();
let mut working = memory.reservation();
working.grow(width_of(capacity * size_of::<i64>() + capacity.div_ceil(8)))?;
let mut slots = vec![0_i64; capacity];
let mut filled = vec![0_u64; capacity.div_ceil(64)];
let mask = capacity - 1;
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;
}
}
}
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;
}
}
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 (mut partition, total) = runs.seed();
let capacity = total.saturating_mul(2).max(64).next_power_of_two();
let mut working = memory.reservation();
let room = total.saturating_sub(partition.rows.len());
working.grow(width_of(
capacity * size_of::<u32>()
+ total * size_of::<CompactNumeric>()
+ room * size_of::<FixedRecord>(),
))?;
let mut buckets = vec![EMPTY_SLOT; capacity];
let mut states: Vec<CompactNumeric> = Vec::with_capacity(total);
let mut overflow = HashMap::new();
partition.rows.reserve(room);
reserving.stop(0);
let timing = stage::Timing::start(Stage::Fold);
let seeded = partition.rows.len();
let all_valid = partition.validity.is_empty();
for source in 0..seeded {
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,
)?;
}
partition.rows.truncate(states.len());
if !all_valid {
partition.validity.truncate(states.len());
}
timing.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();
for (source, &row) in run.rows.iter().enumerate() {
let valid = if all_valid { FixedRecord::ALL } else { run.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.push(row, 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 all_valid = partition.validity.is_empty();
timing.stop(0);
let timing = stage::Timing::start(Stage::Emit);
let mut best: Vec<usize> = Vec::with_capacity(bound.min(states.len()));
for slot in 0..states.len() {
let at = best.partition_point(|&kept| states[kept].count() >= states[slot].count());
if at < bound {
best.insert(at, slot);
best.truncate(bound);
}
}
best.sort_unstable();
let mut output = Vec::with_capacity(best.len());
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(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()?,
]);
}
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,
constants: &[Option<Value>],
group_types: &[LogicalType],
) -> Result<Vec<Chunk>> {
let width = dictionary.len().saturating_add(DENSE_PARTITIONS - 1 - number) / DENSE_PARTITIONS;
let mut dense = vec![0_i64; width];
for run in &partition.runs {
for &code in run {
dense[code as usize / DENSE_PARTITIONS] += 1;
}
}
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 (slot, &count) in dense.iter().enumerate() {
if count == 0 {
continue;
}
codes.push((slot * DENSE_PARTITIONS + number) as u32);
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_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 {
input.div_ceil(65_536).clamp(1, RADIX_PARTITIONS).min(threads.degree())
}
impl Aggregate<'_> {
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() {
if 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 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) {
if 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_pipeline::Sink;
use rudb_plan::{Plan, Slice};
use rudb_vector::{Chunk, Data, Vector};
use super::{
Aggregate, BigIntDistinct, BigIntDistinctRuns, Call, CompactNumeric, Distinct,
EncodedCountPartition, EncodedCountRecord, EncodedCountRuns, FixedPartition, FixedRecord,
FixedRuns, Signed, bigint_distinct_partition, encoded_count_partition, fixed_partition,
};
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 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..5_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,
"five thousand groups on two instances is meant to take the partitioned path"
);
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_under_a_pushed_down_bound_keeps_its_table_in_one_piece() {
let plan = parsed("Aggregate #1 groups=[#0.0::INTEGER] aggregates=[count_star()::BIGINT]");
let (aggregate, out) = aggregate(&plan);
let aggregate = aggregate.top_counts(1_000);
let mut left = aggregate.local();
let mut right = aggregate.local();
let values: Vec<i32> = (0..5_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,
"sixteen partitions of three hundred groups would let the bound through untouched"
);
aggregate.finalize(&rudb_pipeline::Lease::alone()).expect("the answer");
assert_eq!(
answer(&out).len(),
1_000,
"the bound is applied once and not once per partition"
);
}
#[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_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 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 = EncodedCountPartition::default();
partition.push(row(1, 2, 0), EncodedCountRecord::ALL);
partition.push(row(1, 2, 0), EncodedCountRecord::ALL);
partition.push(row(1, 2, 1), EncodedCountRecord::ALL);
partition.push(row(0, 2, 0), EncodedCountRecord::SECOND | EncodedCountRecord::THIRD);
let leading = [LogicalType::BigInt, LogicalType::BigInt];
let part = encoded_count_partition(
&mut EncodedCountRuns { runs: vec![partition] },
&dictionary,
&leading,
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_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 = EncodedCountPartition::default();
narrow.push(row(1, 0), EncodedCountRecord::ALL);
let mut widest = EncodedCountPartition::default();
for _ in 0..3 {
widest.push(row(1, 0), EncodedCountRecord::ALL);
}
widest.push(row(2, 1), EncodedCountRecord::ALL);
let mut late = EncodedCountPartition::default();
late.push(row(0, 0), EncodedCountRecord::SECOND | EncodedCountRecord::THIRD);
late.push(row(1, 0), EncodedCountRecord::ALL);
let leading = [LogicalType::BigInt];
let part = encoded_count_partition(
&mut EncodedCountRuns { runs: vec![narrow, widest, late] },
&dictionary,
&leading,
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_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 = EncodedCountPartition::default();
partition.push(row(1, 0), EncodedCountRecord::ALL);
partition.push(row(1, 0), EncodedCountRecord::ALL);
partition.push(row(1, 1), EncodedCountRecord::ALL);
partition.push(row(0, 0), EncodedCountRecord::SECOND | EncodedCountRecord::THIRD);
let leading = [LogicalType::SmallInt];
let part = encoded_count_partition(
&mut EncodedCountRuns { runs: vec![partition] },
&dictionary,
&leading,
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);
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 = EncodedCountPartition::default();
let row = EncodedCountRecord { first: 0, second: 0, hash: 7, third: 0 };
let some = EncodedCountRecord::SECOND | EncodedCountRecord::THIRD;
encoded.push(row, some);
encoded.push(row, EncodedCountRecord::ALL);
assert_eq!(encoded.validity, vec![some, EncodedCountRecord::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 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);
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 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 = FixedPartition::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,
},
Call {
name: "sum".into(),
args: Vec::new(),
distinct: false,
filter: None,
returns: LogicalType::HugeInt,
affine: None,
},
Call {
name: "avg".into(),
args: Vec::new(),
distinct: false,
filter: None,
returns: LogicalType::Double,
affine: 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);
}
}