use std::cmp::Ordering;
use std::ops::Range;
use std::sync::atomic::{AtomicUsize, Ordering as Atomic};
use std::sync::{Mutex, RwLock};
use rudb_common::{Error, LogicalType, Memory, Reservation, Result, Session, Value};
use rudb_pipeline::{Lease, Progress, Sink};
use rudb_plan::{Plan, Slice, SortKey};
use rudb_vector::{
Buffer, Chunk, Data, VECTOR_SIZE, Validity, Vector, interleave_placed, placed_string_rows,
strings_placeable,
};
use crate::merged::{ORDER, Sorted, order_of, ordering};
use crate::normal::{self, Normal};
use crate::pairs::in_parallel;
use crate::prepared::{Prepared, Scratch};
use crate::rows;
use crate::runs::Runs;
use crate::schema::Schema;
pub(crate) type Sortable = (Vec<Value>, Arrival, Source);
pub(crate) type Arrival = (u64, u64);
pub(crate) type Source = (u32, u32);
pub(crate) type Normalized = (Normal, Arrival, Source);
const BESIDE: u64 = (size_of::<Sortable>() - size_of::<Vec<Value>>()) as u64;
const NORMALIZED: u64 = size_of::<Normalized>() as u64;
#[derive(Debug)]
pub(crate) struct Sort {
keys: Vec<SortKey>,
widths: Option<Vec<usize>>,
ranked: Option<Vec<(usize, bool)>>,
exprs: Prepared,
types: Vec<LogicalType>,
memory: Memory,
gathered: Mutex<Combined>,
charged: Mutex<Vec<Reservation>>,
held: Mutex<Reservation>,
runs: Mutex<Vec<Runs>>,
instances: AtomicUsize,
out: Sorted,
}
#[derive(Debug, Default)]
struct Combined {
chunks: Vec<Chunk>,
rows: Keyed,
sorted: Vec<(u32, Vec<Normalized>)>,
}
impl Combined {
fn empty(rows: &Keyed) -> Self {
Self { chunks: Vec::new(), rows: rows.empty(), sorted: Vec::new() }
}
fn laid(mut self) -> Result<(Vec<Chunk>, Keyed)> {
for (base, run) in self.sorted {
self.rows.absorb(Keyed::Normal(run), base)?;
}
Ok((self.chunks, self.rows))
}
}
#[derive(Debug)]
enum Keyed {
Normal(Vec<Normalized>),
Ranked(Ranked),
Valued(Vec<Sortable>),
}
impl Default for Keyed {
fn default() -> Self {
Self::Valued(Vec::new())
}
}
#[derive(Debug)]
struct Ranked {
layout: Vec<(usize, bool)>,
keys: Vec<Vec<Vector>>,
arrivals: Vec<Arrival>,
rows: Vec<Normalized>,
}
impl Ranked {
fn new(layout: Vec<(usize, bool)>) -> Self {
Self { layout, keys: Vec::new(), arrivals: Vec::new(), rows: Vec::new() }
}
fn sort(&mut self, keys: &[SortKey]) -> Result<()> {
let mut rows: Vec<Normalized> = Vec::with_capacity(self.arrivals.len());
for (chunk, columns) in self.keys.iter().enumerate() {
let chunk = u32::try_from(chunk).map_err(|_| too_many())?;
let len = columns.first().map_or(0, Vector::len);
let first = rows.len();
let arrivals = self.arrivals.get(first..first + len).ok_or_else(mismatched)?;
rows.extend(
arrivals
.iter()
.zip(0..)
.map(|(&arrival, row)| ([0; normal::WIDTH], arrival, (chunk, row))),
);
}
if rows.len() != self.arrivals.len() {
return Err(mismatched());
}
let mut written = 0;
for (position, &(wide, ranked)) in self.layout.iter().enumerate() {
let key = *keys.get(position).ok_or_else(mismatched)?;
let ranks = if ranked { Some(self.ranks(position)?) } else { None };
let mut first = 0;
for columns in &self.keys {
let column = columns.get(position).ok_or_else(mismatched)?;
let len = column.len();
let into = rows.get_mut(first..first + len).ok_or_else(mismatched)?;
let into = into.iter_mut().map(|row| &mut row.0);
match &ranks {
Some(ranks) => {
let ranks = ranks.get(first..first + len).ok_or_else(mismatched)?;
let ranks = Buffer::from_vec(ranks.to_vec());
let valid = Validity::from_iter(len, |at| !column.is_null_at(at));
let ranked = Vector::flat(LogicalType::UInteger, Data::UInt32(ranks))?
.with_validity(valid);
normal::write_column(into, written, wide, &ranked, key)?;
}
None => normal::write_column(into, written, wide, column, key)?,
}
first += len;
}
written += wide;
}
rows.sort_unstable_by(|left, right| {
left.0.cmp(&right.0).then_with(|| left.1.cmp(&right.1))
});
self.rows = rows;
Ok(())
}
fn ranks(&self, position: usize) -> Result<Vec<u32>> {
let mut values: Vec<(&[u8], u32)> = Vec::with_capacity(self.arrivals.len());
let mut row = 0_u32;
for columns in &self.keys {
let column = columns.get(position).ok_or_else(mismatched)?;
for at in 0..column.len() {
match column.bytes_at(at) {
Some(bytes) => values.push((bytes, row)),
None if column.is_null_at(at) => {}
None => return Err(Error::internal("a ranked sort key with no bytes")),
}
row = row.checked_add(1).ok_or_else(too_many)?;
}
}
values.sort_unstable_by(|left, right| left.0.cmp(right.0));
let mut ranks = vec![0_u32; row as usize];
let mut rank = 0_u32;
let mut last: Option<&[u8]> = None;
for (bytes, row) in values {
if last.is_some_and(|last| last != bytes) {
rank += 1;
}
last = Some(bytes);
ranks[row as usize] = rank;
}
Ok(ranks)
}
fn footprint(&self) -> u64 {
let keys: usize = self.keys.iter().flatten().map(Vector::footprint).sum();
let arrivals = self.arrivals.len().saturating_mul(size_of::<Arrival>());
u64::try_from(keys.saturating_add(arrivals)).unwrap_or(u64::MAX)
}
}
impl Keyed {
fn empty(&self) -> Self {
match self {
Self::Normal(_) => Self::Normal(Vec::new()),
Self::Ranked(ranked) => Self::Ranked(Ranked::new(ranked.layout.clone())),
Self::Valued(_) => Self::Valued(Vec::new()),
}
}
fn absorb(&mut self, other: Self, base: u32) -> Result<()> {
match (self, other) {
(Self::Normal(into), Self::Normal(from)) => {
into.extend(from.into_iter().map(|(key, arrival, (chunk, row))| {
(key, arrival, (chunk.saturating_add(base), row))
}));
Ok(())
}
(Self::Valued(into), Self::Valued(from)) => {
into.extend(from.into_iter().map(|(key, arrival, (chunk, row))| {
(key, arrival, (chunk.saturating_add(base), row))
}));
Ok(())
}
(Self::Ranked(into), Self::Ranked(from)) => {
if into.keys.len() != base as usize || !from.rows.is_empty() {
return Err(mismatched());
}
into.keys.extend(from.keys);
into.arrivals.extend(from.arrivals);
Ok(())
}
_ => Err(Error::internal("two instances of one sort holding their keys differently")),
}
}
fn sort(&mut self, keys: &[SortKey]) -> Result<()> {
match self {
Self::Normal(rows) => {
rows.sort_unstable_by(|left, right| {
left.0.cmp(&right.0).then_with(|| left.1.cmp(&right.1))
});
Ok(())
}
Self::Ranked(ranked) => ranked.sort(keys),
Self::Valued(rows) => {
let mut failure: Option<Error> = None;
rows.sort_by(|left, right| settled(keys, left, right, &mut failure));
failure.map_or(Ok(()), Err)
}
}
}
fn sources(&self) -> Box<dyn ExactSizeIterator<Item = Source> + '_> {
match self {
Self::Normal(rows) => Box::new(rows.iter().map(|row| row.2)),
Self::Ranked(ranked) => Box::new(ranked.rows.iter().map(|row| row.2)),
Self::Valued(rows) => Box::new(rows.iter().map(|row| row.2)),
}
}
fn len(&self) -> usize {
match self {
Self::Normal(rows) => rows.len(),
Self::Ranked(ranked) => ranked.arrivals.len(),
Self::Valued(rows) => rows.len(),
}
}
fn orders(&self) -> Result<Vec<[u8; ORDER]>> {
match self {
Self::Normal(rows) => {
Ok(rows.iter().map(|(key, arrival, _)| order_of(key, *arrival)).collect())
}
Self::Ranked(_) | Self::Valued(_) => {
Err(Error::internal("a sort holding its keys as values cannot spill"))
}
}
}
fn footprint(&self) -> u64 {
match self {
Self::Normal(rows) => {
u64::try_from(rows.len()).unwrap_or(u64::MAX).saturating_mul(NORMALIZED)
}
Self::Ranked(ranked) => ranked.footprint(),
Self::Valued(rows) => rows.iter().map(|row| rows::footprint(&row.0) + BESIDE).sum(),
}
}
}
#[derive(Debug)]
pub(crate) struct Gathered {
held: Combined,
scratch: Scratch,
charged: Reservation,
place: Place,
runs: Vec<Runs>,
}
#[derive(Debug, Default)]
pub(crate) struct Place {
pub(crate) morsel: u64,
pub(crate) at: u64,
}
impl Place {
pub(crate) fn of(&self, row: usize) -> Arrival {
(self.morsel, self.at.saturating_add(row as u64))
}
pub(crate) fn past(&mut self, rows: usize) {
self.at = self.at.saturating_add(rows as u64);
}
pub(crate) fn start(&mut self, morsel: u64) {
self.morsel = morsel;
self.at = 0;
}
}
impl Sort {
#[must_use]
pub(crate) fn in_session(mut self, session: &Session) -> Self {
self.exprs = self.exprs.in_session(session);
self
}
pub(crate) fn new(
plan: &Plan,
input: &Schema,
keys: Slice,
memory: &Memory,
) -> Result<(Self, Sorted)> {
let keys = plan.sort_key_list(keys).to_vec();
let exprs: Vec<_> = keys.iter().map(|key| key.expr).collect();
let types: Vec<_> = exprs.iter().map(|&expr| plan.expr_type(expr).clone()).collect();
let out = Sorted::new();
let widths = normal::layout(&types);
let ranked = if widths.is_some() { None } else { normal::ranked_layout(&types) };
let rows = keyed(widths.as_ref(), ranked.as_ref());
let sort = Self {
widths,
ranked,
exprs: Prepared::new(plan, &exprs, input)?,
keys,
types: input.types(),
memory: memory.clone(),
gathered: Mutex::new(Combined { chunks: Vec::new(), rows, sorted: Vec::new() }),
charged: Mutex::new(Vec::new()),
held: Mutex::new(memory.reservation()),
runs: Mutex::new(Vec::new()),
instances: AtomicUsize::new(0),
out: out.clone(),
};
Ok((sort, out))
}
fn tight(&self, local: &Gathered) -> bool {
if self.widths.is_none() {
return false;
}
let Some(limit) = self.memory.limit() else { return false };
if self.memory.used() < limit - limit / 4 {
return false;
}
let instances = self.instances.load(Atomic::Relaxed).max(1) as u64;
let worth = ((limit / 16) / instances).max(1 << 20);
local.charged.bytes() >= worth
}
fn spill(&self, local: &mut Gathered) -> Result<()> {
let empty = Combined::empty(&local.held.rows);
let Combined { chunks, rows, .. } = std::mem::replace(&mut local.held, empty);
let mut charged = vec![std::mem::replace(&mut local.charged, self.memory.reservation())];
let file = self.run(chunks, rows, &mut charged)?;
if let Some(left) = charged.pop() {
local.charged = left;
}
local.runs.push(file);
Ok(())
}
fn run(
&self,
chunks: Vec<Chunk>,
mut rows: Keyed,
charged: &mut Vec<Reservation>,
) -> Result<Runs> {
rows.sort(&self.keys)?;
let total = rows.len();
if u32::try_from(total).is_err() {
return Err(too_many());
}
let orders = rows.orders()?;
let order = order(&chunks, &rows)?;
let taken = rows.footprint();
drop(rows);
give(charged, taken);
let mut types = self.types.clone();
types.push(LogicalType::Blob);
let mut file = Runs::new("sort", types)?;
file.begin(total)?;
lay(&self.types, chunks, &order, total, charged, |whole| file.column(whole))?;
for block in 0..total.div_ceil(VECTOR_SIZE) {
let start = block * VECTOR_SIZE;
let Some(orders) = orders.get(start..(start + VECTOR_SIZE).min(total)) else {
return Err(Error::internal("a sorted run with fewer keys in it than rows"));
};
file.part(&ordering(orders)?)?;
}
Ok(file)
}
}
impl Sink for Sort {
type Local = Gathered;
fn local(&self) -> Gathered {
let rows = keyed(self.widths.as_ref(), self.ranked.as_ref());
self.instances.fetch_add(1, Atomic::Relaxed);
Gathered {
held: Combined { chunks: Vec::new(), rows, sorted: Vec::new() },
scratch: self.exprs.scratch(),
charged: self.memory.reservation(),
place: Place::default(),
runs: Vec::new(),
}
}
fn at(&self, morsel: &rudb_pipeline::Morsel, local: &mut Gathered) -> Result<()> {
local.place.start(morsel.index());
Ok(())
}
fn sink(&self, chunk: &Chunk, local: &mut Gathered) -> Result<Progress> {
if chunk.is_empty() {
return Ok(Progress::More);
}
let mut keys = Vec::with_capacity(self.keys.len());
self.exprs.evaluate(chunk, &mut local.scratch, &mut keys)?;
let at = u32::try_from(local.held.chunks.len()).map_err(|_| too_many())?;
let mut taken = u64::try_from(chunk.footprint()).unwrap_or(u64::MAX);
match (&self.widths, &mut local.held.rows) {
(Some(widths), Keyed::Normal(rows)) => {
let first = rows.len();
let place = &local.place;
rows.extend(
(0..chunk.len())
.map(|row| ([0; normal::WIDTH], place.of(row), (at, row as u32))),
);
let fresh = rows.get_mut(first..).ok_or_else(mismatched)?;
let mut written = 0;
for (position, column) in keys.iter().enumerate() {
let wide = *widths.get(position).ok_or_else(mismatched)?;
let key = *self.keys.get(position).ok_or_else(mismatched)?;
normal::write_column(
fresh.iter_mut().map(|row| &mut row.0),
written,
wide,
column,
key,
)?;
written += wide;
}
taken += NORMALIZED * chunk.len() as u64;
}
(None, Keyed::Ranked(ranked)) => {
taken += keys.iter().map(|column| column.footprint() as u64).sum::<u64>();
taken += (size_of::<Arrival>() * chunk.len()) as u64;
let place = &local.place;
ranked.arrivals.extend((0..chunk.len()).map(|row| place.of(row)));
ranked.keys.push(keys);
}
(None, Keyed::Valued(rows)) => {
for row in 0..chunk.len() {
let key: Vec<Value> = keys
.iter()
.map(|column| column.try_value_at(row))
.collect::<Result<_>>()?;
taken += rows::footprint(&key) + BESIDE;
let row = u32::try_from(row).map_err(|_| too_many())?;
rows.push((key, local.place.of(row as usize), (at, row)));
}
}
_ => return Err(mismatched()),
}
local.held.chunks.push(chunk.clone());
local.place.past(chunk.len());
local.charged.grow(taken)?;
if self.tight(local) {
self.spill(local)?;
}
Ok(Progress::More)
}
fn combine(&self, mut local: Gathered) -> Result<()> {
if let Keyed::Normal(rows) = &mut local.held.rows {
rows.sort_unstable_by(|left, right| {
left.0.cmp(&right.0).then_with(|| left.1.cmp(&right.1))
});
}
self.runs.lock().map_err(poisoned)?.extend(local.runs);
let mut gathered = self.gathered.lock().map_err(poisoned)?;
let base = u32::try_from(gathered.chunks.len()).map_err(|_| too_many())?;
match local.held.rows {
Keyed::Normal(rows) => gathered.sorted.push((base, rows)),
held @ (Keyed::Ranked(_) | Keyed::Valued(_)) => gathered.rows.absorb(held, base)?,
}
gathered.chunks.extend(local.held.chunks);
self.charged.lock().map_err(poisoned)?.push(local.charged);
Ok(())
}
fn finalize(&self, threads: &Lease<'_>) -> Result<()> {
let combined = {
let mut gathered = self.gathered.lock().map_err(poisoned)?;
let empty = Combined::empty(&gathered.rows);
std::mem::replace(&mut *gathered, empty)
};
let mut charged = std::mem::take(&mut *self.charged.lock().map_err(poisoned)?);
let mut files = std::mem::take(&mut *self.runs.lock().map_err(poisoned)?);
if !files.is_empty() {
let (chunks, rows) = combined.laid()?;
if rows.len() > 0 {
files.push(self.run(chunks, rows, &mut charged)?);
}
return self.out.merge(files, self.types.clone());
}
let Combined { chunks, mut rows, sorted } = combined;
let total = rows.len() + sorted.iter().map(|(_, run)| run.len()).sum::<usize>();
if u32::try_from(total).is_err() {
return Err(too_many());
}
let order = match &rows {
Keyed::Normal(_) => {
let runs: Vec<(u32, &[Normalized])> =
sorted.iter().map(|(base, run)| (*base, run.as_slice())).collect();
merged(&runs, total, &starts(&chunks), threads)?
}
Keyed::Ranked(_) | Keyed::Valued(_) => {
rows.sort(&self.keys)?;
order(&chunks, &rows)?
}
};
let taken = rows.footprint()
+ u64::try_from(total - rows.len()).unwrap_or(u64::MAX).saturating_mul(NORMALIZED);
drop(rows);
drop(sorted);
give(&mut charged, taken);
let inverse = placed(&order, threads)?;
let order = Placing { order: &order, inverse: inverse.as_deref() };
let mut held = self.held.lock().map_err(poisoned)?;
let out = gathered(&self.types, chunks, order, total, &mut held, &mut charged, threads)?;
self.out.hold(out)?;
Ok(())
}
}
fn keyed(widths: Option<&Vec<usize>>, ranked: Option<&Vec<(usize, bool)>>) -> Keyed {
match (widths, ranked) {
(Some(_), _) => Keyed::Normal(Vec::new()),
(None, Some(layout)) => Keyed::Ranked(Ranked::new(layout.clone())),
(None, None) => Keyed::Valued(Vec::new()),
}
}
fn order(chunks: &[Chunk], rows: &Keyed) -> Result<Vec<usize>> {
let starts = starts(chunks);
rows.sources()
.map(|(chunk, row)| {
let (Some(&start), Some(len)) =
(starts.get(chunk as usize), chunks.get(chunk as usize).map(Chunk::len))
else {
return Err(Error::internal(
"a sorted row pointing outside the chunks it came from",
));
};
if row as usize >= len {
return Err(Error::internal(
"a sorted row pointing outside the chunks it came from",
));
}
Ok(start + row as usize)
})
.collect()
}
fn starts(chunks: &[Chunk]) -> Vec<usize> {
let mut starts = Vec::with_capacity(chunks.len());
let mut start = 0;
for chunk in chunks {
starts.push(start);
start += chunk.len();
}
starts
}
fn merged(
runs: &[(u32, &[Normalized])],
total: usize,
starts: &[usize],
threads: &Lease<'_>,
) -> Result<Vec<usize>> {
let bases: Vec<u32> =
runs.iter().filter(|(_, run)| !run.is_empty()).map(|&(base, _)| base).collect();
let runs: Vec<&[Normalized]> =
runs.iter().map(|&(_, run)| run).filter(|run| !run.is_empty()).collect();
let at = |base: u32, row: &Normalized| -> Result<usize> {
let (chunk, row) = row.2;
starts
.get(chunk.saturating_add(base) as usize)
.map(|start| start + row as usize)
.ok_or_else(|| Error::internal("a sorted row pointing outside the chunks it came from"))
};
let first = |left: &Normalized, right: &Normalized| {
left.0.cmp(&right.0).then_with(|| left.1.cmp(&right.1))
};
let degree = threads.degree().max(1);
let parts = if runs.len() > 1 { degree * PARTS_A_THREAD } else { 1 };
let mut sample: Vec<&Normalized> = Vec::new();
for run in &runs {
let take = (run.len() * SAMPLE * parts / total.max(1)).clamp(1, run.len());
sample.extend((0..take).filter_map(|index| run.get(index * run.len() / take)));
}
sample.sort_unstable_by(|left, right| first(left, right));
let splitters: Vec<&Normalized> =
(1..parts).filter_map(|part| sample.get(part * sample.len() / parts).copied()).collect();
let cuts: Vec<Vec<usize>> =
runs.iter()
.map(|run| {
let mut cut = Vec::with_capacity(splitters.len() + 2);
cut.push(0);
cut.extend(splitters.iter().map(|splitter| {
run.partition_point(|row| first(row, splitter) == Ordering::Less)
}));
cut.push(run.len());
cut
})
.collect();
let parts = splitters.len() + 1;
let mut out = vec![0usize; total];
let mut slots: Vec<Mutex<&mut [usize]>> = Vec::with_capacity(parts);
let mut rest: &mut [usize] = &mut out;
for part in 0..parts {
let len = cuts.iter().map(|cut| cut[part + 1].saturating_sub(cut[part])).sum();
let (head, tail) = std::mem::take(&mut rest).split_at_mut(len);
slots.push(Mutex::new(head));
rest = tail;
}
in_parallel(threads, parts, degree, "merged sorted part", |part| {
let slot = slots.get(part).ok_or_else(|| Error::internal("a merged part past the end"))?;
let mut into = slot.lock().map_err(poisoned)?;
let pieces: Vec<(u32, &[Normalized])> = runs
.iter()
.zip(&bases)
.zip(&cuts)
.filter_map(|((run, &base), cut)| Some((base, run.get(cut[part]..cut[part + 1])?)))
.filter(|(_, piece)| !piece.is_empty())
.collect();
if let [(base, piece)] = pieces.as_slice() {
for (slot, row) in into.iter_mut().zip(piece.iter()) {
*slot = at(*base, row)?;
}
return Ok(());
}
let mut heads: std::collections::BinaryHeap<Head<'_>> = pieces
.iter()
.enumerate()
.filter_map(|(piece, (_, rows))| rows.first().map(|row| Head { row, piece, index: 0 }))
.collect();
let mut written = 0;
while let Some(Head { row, piece, index }) = heads.pop() {
let slot = into
.get_mut(written)
.ok_or_else(|| Error::internal("a merged part longer than its cut"))?;
let (base, rows) = pieces
.get(piece)
.copied()
.ok_or_else(|| Error::internal("a merged piece past the end"))?;
*slot = at(base, row)?;
written += 1;
if let Some(next) = rows.get(index + 1) {
heads.push(Head { row: next, piece, index: index + 1 });
}
}
Ok(())
})?;
Ok(out)
}
const PARTS_A_THREAD: usize = 4;
const SAMPLE: usize = 64;
struct Head<'a> {
row: &'a Normalized,
piece: usize,
index: usize,
}
impl Ord for Head<'_> {
fn cmp(&self, other: &Self) -> Ordering {
other.row.0.cmp(&self.row.0).then_with(|| other.row.1.cmp(&self.row.1))
}
}
impl PartialOrd for Head<'_> {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl PartialEq for Head<'_> {
fn eq(&self, other: &Self) -> bool {
self.cmp(other) == Ordering::Equal
}
}
impl Eq for Head<'_> {}
fn give(charged: &mut Vec<Reservation>, mut bytes: u64) {
while bytes > 0 {
let Some(last) = charged.last_mut() else { return };
let held = last.bytes();
if held > bytes {
last.shrink(bytes);
return;
}
bytes -= held;
charged.pop();
}
}
fn lay(
types: &[LogicalType],
chunks: Vec<Chunk>,
order: &[usize],
rows: usize,
charged: &mut Vec<Reservation>,
mut each: impl FnMut(&Vector) -> Result<()>,
) -> Result<()> {
if rows == 0 {
return Ok(());
}
for (ty, pieces) in types.iter().zip(transposed(types, chunks)?) {
let (whole, given) = column(ty, pieces, Placing { order, inverse: None })?;
give(charged, given);
each(&whole)?;
}
Ok(())
}
fn transposed(types: &[LogicalType], chunks: Vec<Chunk>) -> Result<Vec<Vec<Vector>>> {
let mut pieces: Vec<Vec<Vector>> = vec![Vec::with_capacity(chunks.len()); types.len()];
for chunk in chunks {
for (position, column) in chunk.into_columns().into_iter().enumerate() {
let Some(into) = pieces.get_mut(position) else {
return Err(Error::internal("a sorted chunk wider than the schema it came from"));
};
into.push(column);
}
}
Ok(pieces)
}
#[derive(Clone, Copy)]
struct Placing<'a> {
order: &'a [usize],
inverse: Option<&'a [u32]>,
}
const PLACED_RUNS: usize = 16 * 1024;
fn placed(order: &[usize], threads: &Lease<'_>) -> Result<Option<Vec<u32>>> {
let runs = 1 + order.windows(2).filter(|pair| pair[1] < pair[0]).count();
if runs > PLACED_RUNS || u32::try_from(order.len()).is_err() {
return Ok(None);
}
let rows = order.len();
let per = rows.div_ceil(threads.degree().max(1)).max(1);
let parts =
in_parallel(threads, rows.div_ceil(per), threads.degree(), "turned order", |part| {
let first = part * per;
let mut out = vec![0u32; per.min(rows - first)];
for (at, &row) in order.iter().enumerate() {
if let Some(slot) = row.checked_sub(first).and_then(|offset| out.get_mut(offset)) {
*slot = at as u32;
}
}
Ok(out)
})?;
Ok(Some(parts.concat()))
}
fn column(ty: &LogicalType, pieces: Vec<Vector>, order: Placing<'_>) -> Result<(Vector, u64)> {
let whole = interleave_placed(ty, &pieces, order.order, order.inverse)?.into_pages();
let given = pieces.iter().map(Vector::footprint).sum::<usize>();
drop(pieces);
Ok((whole, u64::try_from(given).unwrap_or(u64::MAX)))
}
fn split(
types: &[LogicalType],
pieces: &[Vec<Vector>],
longest: &[usize],
order: Placing<'_>,
rows: usize,
degree: usize,
) -> Vec<(usize, Option<Range<usize>>)> {
let mut tasks = Vec::with_capacity(longest.len() + degree);
for &position in longest {
let placeable = order.inverse.is_some()
&& types
.get(position)
.zip(pieces.get(position))
.is_some_and(|(ty, pieces)| strings_placeable(ty, pieces));
let ranges = if placeable {
PLACED_RANGES.min(degree).min(rows / MIN_PLACED_RANGE).max(1)
} else {
1
};
if ranges == 1 {
tasks.push((position, None));
continue;
}
let per = rows.div_ceil(ranges).div_ceil(VECTOR_SIZE) * VECTOR_SIZE;
let mut start = 0;
while start < rows {
let end = (start + per).min(rows);
tasks.push((position, Some(start..end)));
start = end;
}
}
tasks
}
const PLACED_RANGES: usize = 3;
const MIN_PLACED_RANGE: usize = 32 * VECTOR_SIZE;
fn gathered(
types: &[LogicalType],
chunks: Vec<Chunk>,
order: Placing<'_>,
rows: usize,
held: &mut Reservation,
charged: &mut Vec<Reservation>,
threads: &Lease<'_>,
) -> Result<Vec<Chunk>> {
if rows == 0 {
return Ok(Vec::new());
}
let pieces = transposed(types, chunks)?;
let mut longest: Vec<usize> = (0..types.len()).collect();
longest.sort_by_key(|&position| {
let stringy = types
.get(position)
.is_some_and(|ty| matches!(ty, LogicalType::Varchar | LogicalType::Blob));
let bytes = pieces.get(position).map_or(0, |run| run.iter().map(Vector::footprint).sum());
std::cmp::Reverse((stringy, bytes))
});
let tasks = split(types, &pieces, &longest, order, rows, threads.degree());
let mut left = vec![0usize; types.len()];
for &(position, _) in &tasks {
if let Some(count) = left.get_mut(position) {
*count += 1;
}
}
let left: Vec<AtomicUsize> = left.into_iter().map(AtomicUsize::new).collect();
let pieces: Vec<RwLock<Vec<Vector>>> = pieces.into_iter().map(RwLock::new).collect();
let charged = Mutex::new(charged);
let laid = in_parallel(threads, tasks.len(), threads.degree(), "laid sorted column", |task| {
let Some((position, range)) = tasks.get(task).cloned() else {
return Err(Error::internal("a sorted column task past the end"));
};
let (Some(ty), Some(held), Some(left)) =
(types.get(position), pieces.get(position), left.get(position))
else {
return Err(Error::internal("a sorted column past the end of the schema"));
};
let Some(range) = range else {
let pieces = std::mem::take(&mut *held.write().map_err(poisoned)?);
let (whole, given) = column(ty, pieces, order)?;
give(*charged.lock().map_err(poisoned)?, given);
return Ok((position, 0, whole));
};
let start = range.start;
let part = {
let pieces = held.read().map_err(poisoned)?;
let inverse = order.inverse.unwrap_or_default();
placed_string_rows(ty, &pieces, inverse, range)?.into_pages()
};
if left.fetch_sub(1, Atomic::AcqRel) == 1 {
let pieces = std::mem::take(&mut *held.write().map_err(poisoned)?);
let given = pieces.iter().map(Vector::footprint).sum::<usize>();
drop(pieces);
give(*charged.lock().map_err(poisoned)?, u64::try_from(given).unwrap_or(u64::MAX));
}
Ok((position, start, part))
})?;
let mut wholes: Vec<Vec<(usize, Vector)>> = vec![Vec::new(); types.len()];
for (position, start, part) in laid {
if let Some(parts) = wholes.get_mut(position) {
parts.push((start, part));
}
}
let blocks = rows.div_ceil(VECTOR_SIZE);
let mut columns: Vec<Vec<Vector>> = vec![Vec::with_capacity(types.len()); blocks];
for parts in &mut wholes {
parts.sort_unstable_by_key(|&(start, _)| start);
for (block, into) in columns.iter_mut().enumerate() {
let start = block * VECTOR_SIZE;
let Some((from, part)) = parts.iter().rev().find(|&&(from, _)| from <= start) else {
return Err(Error::internal("a sorted column nobody laid"));
};
into.push(part.slice(start - from, (rows - start).min(VECTOR_SIZE))?);
}
}
drop(wholes);
let mut built = Vec::with_capacity(blocks);
for (block, columns) in columns.into_iter().enumerate() {
let start = block * VECTOR_SIZE;
let chunk = Chunk::with_rows(columns, (rows - start).min(VECTOR_SIZE))?;
held.grow(u64::try_from(chunk.footprint()).unwrap_or(u64::MAX))?;
built.push(chunk);
}
Ok(built)
}
fn too_many() -> Error {
Error::internal("a sort of more than 4294967295 rows")
}
fn mismatched() -> Error {
Error::internal("a sort holding its keys two ways at once")
}
pub(crate) fn settled(
keys: &[SortKey],
left: &Sortable,
right: &Sortable,
failure: &mut Option<Error>,
) -> Ordering {
match compare(keys, &left.0, &right.0, failure) {
Ordering::Equal => left.1.cmp(&right.1),
ordering => ordering,
}
}
pub(crate) fn compare(
keys: &[SortKey],
left: &[Value],
right: &[Value],
failure: &mut Option<Error>,
) -> Ordering {
for (at, key) in keys.iter().enumerate() {
let ordering = match rank(&left[at], &right[at], *key) {
Ok(ordering) => ordering,
Err(error) => {
failure.get_or_insert(error);
Ordering::Equal
}
};
if ordering != Ordering::Equal {
return ordering;
}
}
Ordering::Equal
}
pub(crate) fn rank(left: &Value, right: &Value, key: SortKey) -> Result<Ordering> {
match (left.is_null(), right.is_null()) {
(true, true) => Ok(Ordering::Equal),
(true, false) => Ok(if key.nulls_first { Ordering::Less } else { Ordering::Greater }),
(false, true) => Ok(if key.nulls_first { Ordering::Greater } else { Ordering::Less }),
(false, false) => {
let ordering = rudb_kernels::order(left, right)?;
Ok(if key.descending { ordering.reverse() } else { ordering })
}
}
}
fn poisoned<T>(_: T) -> Error {
Error::internal("a thread panicked while holding the rows a sort is gathering")
}
#[cfg(test)]
mod tests {
use rudb_pipeline::{Lease, Pool};
use rudb_vector::VECTOR_SIZE;
use super::{Normalized, merged};
use crate::normal::WIDTH;
fn shuffled(count: usize) -> Vec<(u64, u64)> {
let mut state = 0x9e37_79b9_7f4a_7c15_u64;
(0..count as u64)
.map(|arrival| {
state ^= state << 13;
state ^= state >> 7;
state ^= state << 17;
(state % 97, arrival)
})
.collect()
}
#[test]
fn merged_runs_read_in_the_order_one_sort_would() {
let pool = Pool::new(4);
for threads in [Lease::alone(), pool.lease(4)] {
for lengths in [
vec![],
vec![VECTOR_SIZE * 3 + 1],
vec![5, 0, VECTOR_SIZE, 3 * VECTOR_SIZE + 11, 1],
] {
let total: usize = lengths.iter().sum();
let keys = shuffled(total);
let (mut runs, mut starts) = (Vec::new(), Vec::new());
let mut at = 0;
for (chunk, &length) in lengths.iter().enumerate() {
starts.push(at);
let mut run: Vec<Normalized> = (0..length)
.map(|row| {
let (key, arrival) = keys[at + row];
let mut normal = [0; WIDTH];
normal[..8].copy_from_slice(&key.to_be_bytes());
(normal, (arrival, 0), (0, row as u32))
})
.collect();
run.sort_unstable_by(|left, right| {
left.0.cmp(&right.0).then_with(|| left.1.cmp(&right.1))
});
runs.push((chunk as u32, run));
at += length;
}
let mut expected: Vec<(u64, u64, usize)> = keys
.iter()
.enumerate()
.map(|(index, &(key, arrival))| (key, arrival, index))
.collect();
expected.sort_unstable();
let expected: Vec<usize> =
expected.into_iter().map(|(_, _, index)| index).collect();
let runs: Vec<(u32, &[Normalized])> =
runs.iter().map(|(base, run)| (*base, run.as_slice())).collect();
let got = merged(&runs, total, &starts, &threads).expect("merged");
assert_eq!(got, expected, "runs of {lengths:?} on {} threads", threads.degree());
}
}
}
}