use std::sync::{Arc, Mutex};
use rudb_catalog::Parent;
use rudb_common::{Cancel, Error, LogicalType, Memory, Reservation, Result, Session};
use rudb_graph::Link;
use rudb_pipeline::{Compaction, Gauge, Lease, Progress, Stream, narrow};
use rudb_plan::{ExprRef, JoinKind, Plan};
use rudb_seam::{Context, SeamId, Settings};
use rudb_vector::{Chunk, Data, NO_ROW, Selection, Vector};
use crate::prepared::{Prepared, Scratch};
use crate::register::compaction;
use crate::schema::Schema;
#[derive(Debug)]
pub(crate) struct LinkJoin {
kind: JoinKind,
link: Arc<Link>,
parent: Arc<Parent>,
projected: Vec<(usize, LogicalType)>,
rid: Prepared,
schema: Schema,
compaction: &'static dyn Compaction,
held: Mutex<Reservation>,
cancel: Cancel,
}
#[derive(Debug)]
pub(crate) struct Linking {
scratch: Scratch,
rids: Vec<u32>,
gauge: Gauge,
}
impl LinkJoin {
#[must_use]
pub(crate) fn in_session(mut self, session: &Session) -> Self {
self.rid = self.rid.in_session(session);
self
}
#[expect(clippy::too_many_arguments, reason = "an operator's inputs, none of them a group")]
pub(crate) fn new(
plan: &Plan,
kind: JoinKind,
link: Arc<Link>,
parent: Arc<Parent>,
projected: Vec<(usize, LogicalType)>,
rid: ExprRef,
child: &Schema,
gathered: &Schema,
seams: &Settings,
memory: &Memory,
cancel: Cancel,
) -> Result<Self> {
if !matches!(kind, JoinKind::Inner | JoinKind::Left | JoinKind::Semi | JoinKind::Anti) {
return Err(Error::internal("a link join was built for a kind a link cannot answer"));
}
let types = child.types();
let context = Context::new(SeamId::ChunkCompaction, seams).with_types(&types);
let compaction = compaction().choose(&context)?.strategy();
let rid = Prepared::one(plan, rid, child)?;
Ok(Self {
kind,
link,
parent,
projected,
rid,
schema: Schema::concat(child, gathered),
compaction,
held: Mutex::new(memory.reservation()),
cancel,
})
}
pub(crate) fn schema(&self) -> &Schema {
&self.schema
}
fn resolve(&self, chunk: &Chunk, local: &mut Linking) -> Result<()> {
let rows = chunk.len();
let ids = self.rid.evaluate_one(chunk, &mut local.scratch)?.flatten()?;
let held: &[i64] = match ids.data() {
Some(Data::Int64(values)) if !ids.validity().has_nulls(rows) => values.as_slice(),
_ => return Err(Error::internal("a link join was handed invalid child row ids")),
};
let held = held.get(..rows).ok_or_else(|| {
Error::internal("a link join was handed fewer row ids than the chunk has rows")
})?;
local.rids.clear();
local.rids.reserve(rows);
for &id in held {
let child = u64::try_from(id)
.map_err(|_| Error::internal("a link join was handed a negative child row id"))?;
local.rids.push(match self.link.forward(child) {
Some(parent) => {
u32::try_from(parent).ok().filter(|&rid| rid != NO_ROW).ok_or_else(|| {
Error::internal("a link answered a parent row id a gather cannot hold")
})?
}
None => NO_ROW,
});
}
Ok(())
}
fn gather(&self, chunk: &mut Chunk, rids: &[u32]) -> Result<()> {
if self.projected.is_empty() {
return Ok(());
}
let rows = chunk.len();
let ids = Arc::new(rids.to_vec());
let mut columns: Vec<Vector> = chunk.columns().to_vec();
for (column, ty) in &self.projected {
let source = self.parent.column(*column, ty)?.ok_or_else(|| {
Error::out_of_memory(
"a link join could not hold the parent columns it gathers from".to_string(),
)
})?;
columns.push(Vector::gathered(source, Arc::clone(&ids))?);
}
*chunk = Chunk::with_rows(columns, rows)?;
Ok(())
}
fn read_parent(&self) -> Result<()> {
for (column, ty) in &self.projected {
self.cancel.check()?;
if self.parent.column(*column, ty)?.is_none() {
return Err(Error::out_of_memory(
"a link join could not hold the parent columns it gathers from".to_string(),
));
}
}
let footprint = u64::try_from(self.parent.footprint()).unwrap_or(u64::MAX);
let mut held = self.held.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
let already = held.bytes();
if footprint > already {
held.grow(footprint - already)?;
}
Ok(())
}
}
impl Stream for LinkJoin {
type Local = Linking;
fn local(&self) -> Linking {
Linking { scratch: self.rid.scratch(), rids: Vec::new(), gauge: Gauge::new(1) }
}
fn prepare(&self, _threads: &Lease<'_>) -> Result<()> {
self.read_parent()
}
fn push(&self, chunk: &mut Chunk, local: &mut Linking) -> Result<Progress> {
self.cancel.check()?;
let rows = chunk.len();
if rows == 0 {
*chunk = Chunk::empty(&self.schema.types());
return Ok(Progress::More);
}
self.resolve(chunk, local)?;
match self.kind {
JoinKind::Semi | JoinKind::Anti => {
let hit = self.kind == JoinKind::Semi;
let kept =
Selection::from_predicate(rows, |row| (local.rids[row] != NO_ROW) == hit);
if kept.len() != rows {
narrow(self.compaction, chunk, &kept, &mut local.gauge)?;
}
Ok(Progress::More)
}
JoinKind::Inner => {
let kept = Selection::from_predicate(rows, |row| local.rids[row] != NO_ROW);
if kept.len() != rows {
local.rids = kept.iter().map(|row| local.rids[row]).collect();
narrow(self.compaction, chunk, &kept, &mut local.gauge)?;
}
let rids = std::mem::take(&mut local.rids);
let result = self.gather(chunk, &rids);
local.rids = rids;
result?;
Ok(Progress::More)
}
JoinKind::Left => {
let rids = std::mem::take(&mut local.rids);
let result = self.gather(chunk, &rids);
local.rids = rids;
result?;
Ok(Progress::More)
}
JoinKind::Right
| JoinKind::Full
| JoinKind::Mark
| JoinKind::Single
| JoinKind::Positional => {
Err(Error::internal("a link join was asked for a kind a link cannot answer"))
}
}
}
fn weight(&self) -> usize {
1
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use rudb_catalog::{Parent, Rows};
use rudb_common::{Cancel, Field, LogicalType, Memory, Session, Value};
use rudb_graph::{Link, NO_PARENT};
use rudb_pipeline::Stream;
use rudb_plan::{ColumnBinding, Expr, JoinKind, Plan};
use rudb_seam::Settings;
use rudb_storage::MemoryTable;
use rudb_vector::{Chunk, Vector};
use super::LinkJoin;
use crate::schema::Schema;
fn parent(rows: i32) -> Arc<Parent> {
let mut table = MemoryTable::new(vec![LogicalType::Integer]);
let held: Vec<Value> = (0..rows).map(Value::Integer).collect();
let column = Vector::from_values(LogicalType::Integer, &held).expect("a column");
table.append(Chunk::new(vec![column]).expect("a chunk")).expect("appended");
Arc::new(Parent::new(Rows::Memory(table), 64 * 1024 * 1024))
}
fn child_schema() -> Schema {
Schema::numbered(
vec![
Field::new("l_price", LogicalType::Integer),
Field::new("file_row_number", LogicalType::BigInt),
],
0,
)
}
fn gathered_schema(parent_columns: bool) -> Schema {
let fields = if parent_columns {
vec![Field::new("o_key", LogicalType::Integer)]
} else {
Vec::new()
};
Schema::numbered(fields, 1)
}
fn child(rows: usize) -> Chunk {
let prices: Vec<Value> = (0..rows)
.map(|row| Value::Integer(100 + i32::try_from(row).expect("a small row count")))
.collect();
let prices = Vector::from_values(LogicalType::Integer, &prices).expect("prices");
let rids = Vector::sequence(0, 1, rows);
Chunk::new(vec![prices, rids]).expect("a child chunk")
}
fn link(parents: &[Option<u64>]) -> Arc<Link> {
let held: Vec<u64> = parents.iter().map(|parent| parent.unwrap_or(NO_PARENT)).collect();
let highest = held.iter().filter(|&&p| p != NO_PARENT).max().copied().unwrap_or(0);
Arc::new(Link::build(&held, highest + 1).expect("a link"))
}
fn operator(kind: JoinKind, parents: &[Option<u64>], rows: i32) -> LinkJoin {
let mut plan = Plan::new();
let rid = plan.add_expr(Expr::Column(ColumnBinding::new(0, 1)), LogicalType::BigInt);
let gathers = !matches!(kind, JoinKind::Semi | JoinKind::Anti);
let projected = if gathers { vec![(0, LogicalType::Integer)] } else { Vec::new() };
LinkJoin::new(
&plan,
kind,
link(parents),
parent(rows),
projected,
rid,
&child_schema(),
&gathered_schema(gathers),
&Settings::default(),
&Memory::unlimited(),
Cancel::new(),
)
.expect("the operator is buildable")
.in_session(&Session::default())
}
fn run(operator: &LinkJoin, mut chunk: Chunk) -> Vec<Vec<Value>> {
let mut local = operator.local();
operator.push(&mut chunk, &mut local).expect("the push");
let chunk = chunk.flatten().expect("flattened");
(0..chunk.len()).map(|row| chunk.row(row).collect()).collect()
}
#[test]
fn an_inner_link_join_puts_each_childs_parent_beside_it() {
let operator = operator(JoinKind::Inner, &[Some(2), Some(0), Some(2), Some(1)], 3);
let rows = run(&operator, child(4));
assert_eq!(
rows,
vec![
vec![Value::Integer(100), Value::BigInt(0), Value::Integer(2)],
vec![Value::Integer(101), Value::BigInt(1), Value::Integer(0)],
vec![Value::Integer(102), Value::BigInt(2), Value::Integer(2)],
vec![Value::Integer(103), Value::BigInt(3), Value::Integer(1)],
]
);
}
#[test]
fn an_inner_link_join_drops_the_children_with_no_parent() {
let operator = operator(JoinKind::Inner, &[None, Some(1), None, Some(0)], 2);
let rows = run(&operator, child(4));
assert_eq!(
rows,
vec![
vec![Value::Integer(101), Value::BigInt(1), Value::Integer(1)],
vec![Value::Integer(103), Value::BigInt(3), Value::Integer(0)],
]
);
}
#[test]
fn a_left_link_join_keeps_the_children_with_no_parent_and_gathers_null() {
let operator = operator(JoinKind::Left, &[None, Some(1), None, Some(0)], 2);
let rows = run(&operator, child(4));
assert_eq!(
rows,
vec![
vec![Value::Integer(100), Value::BigInt(0), Value::Null],
vec![Value::Integer(101), Value::BigInt(1), Value::Integer(1)],
vec![Value::Integer(102), Value::BigInt(2), Value::Null],
vec![Value::Integer(103), Value::BigInt(3), Value::Integer(0)],
]
);
}
#[test]
fn a_semi_link_join_keeps_the_children_that_have_a_parent_and_reads_nothing() {
let operator = operator(JoinKind::Semi, &[None, Some(1), None, Some(0)], 2);
let rows = run(&operator, child(4));
assert_eq!(
rows,
vec![
vec![Value::Integer(101), Value::BigInt(1)],
vec![Value::Integer(103), Value::BigInt(3)],
]
);
}
#[test]
fn an_anti_link_join_keeps_the_children_that_have_none() {
let operator = operator(JoinKind::Anti, &[None, Some(1), None, Some(0)], 2);
let rows = run(&operator, child(4));
assert_eq!(
rows,
vec![
vec![Value::Integer(100), Value::BigInt(0)],
vec![Value::Integer(102), Value::BigInt(2)],
]
);
}
#[test]
fn the_parent_is_pointed_at_rather_than_copied() {
let operator = operator(JoinKind::Inner, &[Some(0); 8], 4);
let mut chunk = child(8);
let mut local = operator.local();
operator.push(&mut chunk, &mut local).expect("the push");
let gathered = chunk.column(2).expect("the parent column");
let (source, rids) = gathered.gathered_parts().expect("it is a gather");
assert_eq!(source.len(), 4, "the source is the whole parent column");
assert_eq!(rids.len(), 8, "one id per child row");
assert!(
gathered.footprint() < source.footprint() + 8 * 4 + 64,
"the parent was copied rather than pointed at"
);
}
#[test]
fn a_child_row_id_that_is_not_a_row_id_is_refused() {
let operator = operator(JoinKind::Inner, &[Some(0), Some(0)], 2);
let prices =
Vector::from_values(LogicalType::Integer, &[Value::Integer(1), Value::Integer(2)])
.expect("prices");
let rids = Vector::from_values(LogicalType::BigInt, &[Value::BigInt(-1), Value::BigInt(0)])
.expect("ids");
let mut chunk = Chunk::new(vec![prices, rids]).expect("a chunk");
let mut local = operator.local();
let message = operator.push(&mut chunk, &mut local).unwrap_err().to_string();
assert!(message.contains("negative child row id"), "unhelpful message: {message}");
}
#[test]
fn a_child_past_the_end_of_the_link_has_no_parent() {
let operator = operator(JoinKind::Left, &[Some(0), Some(0)], 2);
let rows = run(&operator, child(4));
assert_eq!(rows.len(), 4);
assert_eq!(rows[2][2], Value::Null, "a child the link does not cover has no parent");
assert_eq!(rows[3][2], Value::Null);
}
}