use std::borrow::Cow;
use std::cmp::Ordering;
use std::sync::{Arc, Mutex};
use rudb_common::bounds::Bound;
use rudb_common::{Error, LogicalType, Memory, Reservation, Result, Session, Value};
use rudb_kernels::{
Comparison, compare as compare_vectors, rank_at, rank_within, select_against_rank, selection,
};
use rudb_pipeline::{Lease, Progress, Sink};
use rudb_plan::{Plan, Slice, SortKey};
use rudb_vector::{Chunk, Selection, Vector};
use crate::buffer::Buffered;
use crate::cutoff::Cutoff;
use crate::prepared::{Prepared, Scratch};
use crate::rows;
use crate::schema::Schema;
use crate::sort::{Place, rank};
const SORTED_BOUND: usize = 64;
#[derive(Debug)]
struct Candidate {
key: Vec<Cell>,
values: Vec<Cell>,
arrival: crate::sort::Arrival,
}
#[derive(Debug, Clone)]
enum Cell {
Ready(Value),
Coded { dictionary: Arc<Vector>, code: u32, rank: Option<u32> },
}
impl Cell {
fn of(column: &Vector, row: usize) -> Result<Self> {
if let Some((codes, dictionary)) = column.shared_dictionary_parts() {
if column.validity().is_valid(row) {
if let Some(code) = codes.get(row) {
let dictionary = Arc::clone(dictionary);
return Ok(Self::Coded { dictionary, code: *code, rank: None });
}
}
}
Ok(Self::Ready(column.try_value_at(row)?))
}
fn keyed(column: &Vector, row: usize) -> Result<Self> {
match placed(column, row) {
Some(cell) => Ok(cell),
None => Ok(Self::Ready(column.try_value_at(row)?)),
}
}
fn value(self) -> Result<Value> {
match self {
Self::Ready(value) => Ok(value),
Self::Coded { dictionary, code, .. } => dictionary.try_value_at(code as usize),
}
}
fn read(&self) -> Result<Cow<'_, Value>> {
match self {
Self::Ready(value) => Ok(Cow::Borrowed(value)),
Self::Coded { dictionary, code, .. } => {
dictionary.try_value_at(*code as usize).map(Cow::Owned)
}
}
}
fn footprint(&self) -> usize {
match self {
Self::Ready(value) => value.footprint(),
Self::Coded { .. } => 0,
}
}
}
fn placed(column: &Vector, row: usize) -> Option<Cell> {
let (dictionary, rank) = rank_at(column, row)?;
let (codes, _) = column.shared_dictionary_parts()?;
let code = *codes.get(row)?;
Some(Cell::Coded { dictionary, code, rank: Some(rank) })
}
fn charge(values: &[Cell]) -> u64 {
let bytes = size_of::<Vec<Cell>>()
+ values.iter().map(Cell::footprint).sum::<usize>()
+ size_of_val(values)
+ usize::try_from(rudb_common::ALLOCATION).unwrap_or(0);
u64::try_from(bytes).unwrap_or(u64::MAX)
}
fn settled(
keys: &[SortKey],
left: &Candidate,
right: &Candidate,
failure: &mut Option<Error>,
) -> Ordering {
match compare(keys, &left.key, &right.key, failure) {
Ordering::Equal => left.arrival.cmp(&right.arrival),
ordering => ordering,
}
}
fn compare(
keys: &[SortKey],
left: &[Cell],
right: &[Cell],
failure: &mut Option<Error>,
) -> Ordering {
for (at, key) in keys.iter().enumerate() {
let ordering = match place(&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
}
fn place(left: &Cell, right: &Cell, key: SortKey) -> Result<Ordering> {
match (left, right) {
(Cell::Ready(here), Cell::Ready(there)) => rank(here, there, key),
(
Cell::Coded { dictionary: one, rank: Some(here), .. },
Cell::Coded { dictionary: other, rank: Some(there), .. },
) if Arc::ptr_eq(one, other) => {
let ordering = here.cmp(there);
Ok(if key.descending { ordering.reverse() } else { ordering })
}
_ => {
let (here, there) = (left.read()?, right.read()?);
rank(&here, &there, key)
}
}
}
#[derive(Debug)]
pub(crate) struct TopN {
keys: Vec<SortKey>,
exprs: Prepared,
types: Vec<LogicalType>,
count: usize,
offset: usize,
bound: usize,
memory: Memory,
rows: Mutex<Vec<Candidate>>,
charged: Mutex<Vec<Reservation>>,
held: Mutex<Reservation>,
cutoff: Option<Arc<Cutoff>>,
out: Buffered,
}
#[derive(Debug)]
pub(crate) struct Running {
kept: Vec<Candidate>,
scratch: Scratch,
charged: Reservation,
failure: Option<Error>,
place: Place,
cut: Option<Vec<Cell>>,
moved: bool,
}
impl TopN {
#[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,
count: u64,
offset: u64,
memory: &Memory,
) -> Result<(Self, Buffered)> {
let count = usize::try_from(count).unwrap_or(usize::MAX);
let offset = usize::try_from(offset).unwrap_or(usize::MAX);
let keys = plan.sort_key_list(keys).to_vec();
let exprs: Vec<_> = keys.iter().map(|key| key.expr).collect();
let out = Buffered::new();
let top = Self {
exprs: Prepared::new(plan, &exprs, input)?,
keys,
types: input.types(),
count,
offset,
bound: count.saturating_add(offset),
memory: memory.clone(),
rows: Mutex::new(Vec::new()),
charged: Mutex::new(Vec::new()),
held: Mutex::new(memory.reservation()),
cutoff: None,
out: out.clone(),
};
Ok((top, out))
}
#[must_use]
pub(crate) fn telling(mut self, cutoff: Arc<Cutoff>) -> Self {
self.cutoff = Some(cutoff);
self
}
fn reached(&self, worst: &[Cell]) {
let Some(cutoff) = self.cutoff.as_ref().filter(|cutoff| cutoff.armed()) else { return };
let Some(first) = worst.first() else { return };
let Ok(value) = first.read() else { return };
let Some(bound) = Bound::of_value(&value) else { return };
cutoff.reached(bound);
}
fn offer(
&self,
keys: &[Vector],
chunk: &Chunk,
row: usize,
local: &mut Running,
trimmed: &mut bool,
) -> Result<u64> {
let Running { kept, failure, place, cut, .. } = local;
if kept.len() > self.bound.saturating_mul(2) {
trim(&self.keys, kept, self.bound, failure);
*trimmed = true;
*cut = (kept.len() == self.bound && self.bound > 0)
.then(|| kept[self.bound - 1].key.clone());
if let Some(reached) = cut.as_ref() {
self.reached(reached);
}
}
let lost = cut
.as_ref()
.is_some_and(|cut| against(&self.keys, keys, row, cut, failure) != Ordering::Less);
if lost {
return Ok(0);
}
let arrival = place.of(row);
hold(keys, chunk, row, arrival, kept)
}
}
impl Sink for TopN {
type Local = Running;
fn local(&self) -> Running {
Running {
kept: Vec::new(),
scratch: self.exprs.scratch(),
charged: self.memory.reservation(),
failure: None,
place: Place::default(),
cut: None,
moved: false,
}
}
fn at(&self, morsel: &rudb_pipeline::Morsel, local: &mut Running) -> Result<()> {
local.place.start(morsel.index());
Ok(())
}
fn sink(&self, chunk: &Chunk, local: &mut Running) -> Result<Progress> {
let mut keys = Vec::with_capacity(self.keys.len());
self.exprs.evaluate(chunk, &mut local.scratch, &mut keys)?;
if self.bound <= SORTED_BOUND {
let rows = chunk.len();
let full = self.bound > 0 && local.kept.len() == self.bound;
let narrowed = full
.then(|| {
let worst = &local.kept[self.bound - 1].key;
let ranked = beats_rank(&self.keys, &keys, worst, rows);
ranked.or_else(|| worth_looking_at(&self.keys, &keys, worst, rows))
})
.flatten();
let offer = |row: usize, local: &mut Running| {
let arrival = local.place.of(row);
keep(
Where { keys: &self.keys, columns: &keys, chunk, row, arrival },
local,
self.bound,
);
};
match narrowed {
Some(kept) => {
for row in kept.iter() {
offer(row, local);
}
}
None => {
for row in 0..rows {
offer(row, local);
}
}
}
local.place.past(chunk.len());
recharge(&local.kept, &mut local.charged)?;
if local.moved && self.bound > 0 && local.kept.len() == self.bound {
local.moved = false;
self.reached(&local.kept[self.bound - 1].key);
}
return Ok(Progress::More);
}
let narrowed = local
.cut
.as_ref()
.and_then(|cut| worth_looking_at(&self.keys, &keys, cut, chunk.len()));
let mut taken = 0;
let mut trimmed = false;
match narrowed {
Some(rows) => {
for row in rows.iter() {
taken += self.offer(&keys, chunk, row, local, &mut trimmed)?;
}
}
None => {
for row in 0..chunk.len() {
taken += self.offer(&keys, chunk, row, local, &mut trimmed)?;
}
}
}
local.place.past(chunk.len());
local.charged.grow(taken)?;
if trimmed {
recharge(&local.kept, &mut local.charged)?;
}
Ok(Progress::More)
}
fn combine(&self, mut local: Running) -> Result<()> {
if let Some(error) = local.failure {
return Err(error);
}
let mut rows = self.rows.lock().map_err(poisoned)?;
rows.extend(local.kept);
let mut failure = None;
trim(&self.keys, &mut rows, self.bound, &mut failure);
if let Some(error) = failure {
return Err(error);
}
recharge(&rows, &mut local.charged)?;
self.charged.lock().map_err(poisoned)?.push(local.charged);
Ok(())
}
fn finalize(&self, _threads: &Lease<'_>) -> Result<()> {
let kept = std::mem::take(&mut *self.rows.lock().map_err(poisoned)?);
let wanted = kept.into_iter().skip(self.offset).take(self.count);
let ordered: Vec<Vec<Value>> = wanted
.map(|candidate| candidate.values.into_iter().map(Cell::value).collect())
.collect::<Result<_>>()?;
let mut held = self.held.lock().map_err(poisoned)?;
let chunks = rows::chunks(&self.types, &ordered, &mut held)?;
self.out.fill(chunks)?;
self.charged.lock().map_err(poisoned)?.clear();
Ok(())
}
}
fn poisoned<T>(_: T) -> Error {
Error::internal("a thread panicked while holding the rows a top N is keeping")
}
fn trim(keys: &[SortKey], kept: &mut Vec<Candidate>, bound: usize, failure: &mut Option<Error>) {
kept.sort_by(|left, right| settled(keys, left, right, failure));
kept.truncate(bound);
}
fn hold(
keys: &[Vector],
chunk: &Chunk,
row: usize,
arrival: crate::sort::Arrival,
kept: &mut Vec<Candidate>,
) -> Result<u64> {
let key: Vec<Cell> =
keys.iter().map(|column| Cell::keyed(column, row)).collect::<Result<_>>()?;
let values: Vec<Cell> =
chunk.columns().iter().map(|column| Cell::of(column, row)).collect::<Result<_>>()?;
let taken = charge(&key) + charge(&values);
kept.push(Candidate { key, values, arrival });
Ok(taken)
}
fn keep(
Where { keys, columns, chunk, row, arrival }: Where<'_>,
local: &mut Running,
bound: usize,
) {
if bound == 0 {
return;
}
let failure = &mut local.failure;
if local.kept.len() == bound
&& against(keys, columns, row, &local.kept[bound - 1].key, failure) != Ordering::Less
{
return;
}
let key: Vec<Cell> = match columns.iter().map(|column| Cell::keyed(column, row)).collect() {
Ok(key) => key,
Err(error) => {
failure.get_or_insert(error);
return;
}
};
let values: Vec<Cell> =
match chunk.columns().iter().map(|column| Cell::of(column, row)).collect() {
Ok(values) => values,
Err(error) => {
failure.get_or_insert(error);
return;
}
};
let at = local.kept.partition_point(|candidate| {
compare(keys, &candidate.key, &key, failure) != Ordering::Greater
});
local.kept.insert(at, Candidate { key, values, arrival });
local.kept.truncate(bound);
local.moved = true;
}
struct Where<'a> {
keys: &'a [SortKey],
columns: &'a [Vector],
chunk: &'a Chunk,
row: usize,
arrival: crate::sort::Arrival,
}
fn against(
keys: &[SortKey],
columns: &[Vector],
row: usize,
held: &[Cell],
failure: &mut Option<Error>,
) -> Ordering {
for (at, key) in keys.iter().enumerate() {
let ordering = match at_row(&columns[at], row, &held[at], *key) {
Ok(ordering) => ordering,
Err(error) => {
failure.get_or_insert(error);
Ordering::Equal
}
};
if ordering != Ordering::Equal {
return ordering;
}
}
Ordering::Equal
}
fn at_row(column: &Vector, row: usize, held: &Cell, key: SortKey) -> Result<Ordering> {
let (dictionary, code, place) = match held {
Cell::Ready(there) => return rank(&column.try_value_at(row)?, there, key),
Cell::Coded { dictionary, code, rank } => (dictionary, code, rank),
};
if let Some(there) = place {
if let Some(here) = rank_within(column, row, dictionary) {
let ordering = here.cmp(there);
return Ok(if key.descending { ordering.reverse() } else { ordering });
}
}
let (here, there) = (column.try_value_at(row)?, dictionary.try_value_at(*code as usize)?);
rank(&here, &there, key)
}
fn worth_looking_at(
keys: &[SortKey],
columns: &[Vector],
worst: &[Cell],
rows: usize,
) -> Option<Selection> {
let key = *keys.first()?;
let bound = worst.first()?.read().ok()?;
let column = columns.first()?;
if bound.is_null() || (key.nulls_first && column.validity().has_nulls(rows)) {
return None;
}
let op = still_wanted(key, keys.len() == 1);
let against = Vector::constant(column.logical_type().clone(), bound.into_owned(), rows);
let flags = compare_vectors(op, column, &against).ok()?;
Some(selection(&flags, rows))
}
fn beats_rank(
keys: &[SortKey],
columns: &[Vector],
worst: &[Cell],
rows: usize,
) -> Option<Selection> {
let key = *keys.first()?;
let column = columns.first()?;
let Cell::Coded { dictionary, rank: Some(rank), .. } = worst.first()? else { return None };
if key.nulls_first && column.validity().has_nulls(rows) {
return None;
}
select_against_rank(still_wanted(key, keys.len() == 1), column, dictionary, *rank, rows)
}
fn still_wanted(key: SortKey, single: bool) -> Comparison {
match (key.descending, single) {
(false, true) => Comparison::Less,
(false, false) => Comparison::LessOrEqual,
(true, true) => Comparison::Greater,
(true, false) => Comparison::GreaterOrEqual,
}
}
fn recharge(kept: &[Candidate], scratch: &mut Reservation) -> Result<()> {
let footprint =
kept.iter().map(|candidate| charge(&candidate.key) + charge(&candidate.values)).sum();
scratch.release();
scratch.grow(footprint)
}