use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use rudb_common::{
Cancel, Error, LogicalType, Memory, Reservation, Result, Session, SessionTimeZone, Value,
};
use rudb_kernels::{Connective, combine, is_true};
use rudb_pipeline::{Lease, Progress, Sink, Stream};
use rudb_plan::{ColumnBinding, CompareOp, Expr, ExprRef, JoinKind, Plan, Slice};
use rudb_vector::{Chunk, Data, VECTOR_SIZE, Validity, Vector};
use crate::buffer::Buffered;
use crate::expr::evaluate_all_in_time_zone;
use crate::gather::{self, Gathering};
use crate::lookup::{Lookup, MISS, Scratch};
use crate::rows;
use crate::schema::Schema;
use crate::side::{Build, PAD, laid_out};
#[derive(Debug)]
pub(crate) struct Join<'a> {
plan: &'a Plan,
kind: JoinKind,
conditions: Vec<ExprRef>,
marker: Option<usize>,
left_schema: Schema,
right_schema: Schema,
combined: Schema,
schema: Schema,
swapped: bool,
memory: Memory,
cancel: Cancel,
right: Buffered,
left: Mutex<Vec<Vec<Value>>>,
charged: Mutex<Vec<Reservation>>,
held: Mutex<Reservation>,
out: Buffered,
time_zone: SessionTimeZone,
}
pub(crate) struct Gathered<'s> {
pub(crate) schema: &'s Schema,
pub(crate) chunks: Buffered,
pub(crate) marker: Option<usize>,
pub(crate) swapped: bool,
}
impl<'a> Join<'a> {
pub(crate) fn new(
plan: &'a Plan,
left: &Schema,
right: Gathered<'_>,
kind: JoinKind,
conditions: Slice,
cancel: &Cancel,
memory: &Memory,
) -> (Self, Buffered) {
let swapped = right.swapped;
let combined = Schema::concat(left, right.schema);
let schema = match kind {
JoinKind::Semi | JoinKind::Anti => left.clone(),
_ if swapped => Schema::concat(right.schema, left),
_ => combined.clone(),
};
let out = Buffered::new();
let marker = if kind == JoinKind::Mark {
right.marker.or_else(|| right.schema.bindings().len().checked_sub(1))
} else {
None
};
let join = Self {
plan,
kind,
conditions: plan.expr_list(conditions).to_vec(),
marker,
left_schema: left.clone(),
right_schema: right.schema.clone(),
combined,
schema,
swapped,
memory: memory.clone(),
cancel: cancel.clone(),
right: right.chunks,
left: Mutex::new(Vec::new()),
charged: Mutex::new(Vec::new()),
held: Mutex::new(memory.reservation()),
out: out.clone(),
time_zone: SessionTimeZone::default(),
};
(join, out)
}
#[must_use]
pub(crate) fn in_session(mut self, session: &Session) -> Self {
self.time_zone = session.session_time_zone();
self
}
pub(crate) fn schema(&self) -> &Schema {
&self.schema
}
fn joined(
&self,
left_rows: &[Vec<Value>],
right_chunks: &[Chunk],
threads: &Lease<'_>,
) -> Result<Vec<Vec<Value>>> {
let mut scratch = self.memory.reservation();
let left_types = self.left_schema.types();
let right_types = self.right_schema.types();
let right = Build::new(&right_types, right_chunks, threads)?;
scratch.grow(right.footprint())?;
let right_rows = right.rows();
if self.kind == JoinKind::Positional {
let rows: Vec<Vec<Value>> = (0..right_rows).map(|at| right.row(at as u32)).collect();
return Ok(positional(left_rows, &rows, left_types.len(), right_types.len()));
}
let marks = match self.kind {
JoinKind::Mark => self.marks(left_rows, right_chunks, &mut scratch)?,
_ => None,
};
let equalities = if self.kind == JoinKind::Mark {
None
} else {
equalities(self.plan, &self.conditions, &self.left_schema, &self.right_schema)
};
let index = match &equalities {
Some(equalities) => Some(lookup(
equalities.gathered(self.plan, &self.right_schema, self.time_zone),
right_chunks,
&self.cancel,
&Lease::alone(),
&mut scratch,
)?),
None => None,
};
let left_slots = match (&equalities, &index) {
(Some(equalities), Some(index)) => found(
equalities.driving(self.plan, &self.left_schema, self.time_zone),
index,
left_rows,
&self.cancel,
&mut scratch,
)?,
_ => Vec::new(),
};
let residual = equalities.as_ref().map(|equalities| Residual {
plan: self.plan,
exprs: &equalities.residual,
combined: &self.combined,
left_types: &left_types,
time_zone: self.time_zone,
});
let scanned_over: &[Chunk] = match index {
Some(_) => &[],
None => right_chunks,
};
let mut matched = vec![false; right_rows];
scratch.grow(u64::try_from(right_rows).unwrap_or(u64::MAX))?;
let mut out: Vec<Vec<Value>> = Vec::new();
let mut kept: Vec<u32> = Vec::new();
let mut chain: Vec<u32> = Vec::new();
for (position, left_row) in left_rows.iter().enumerate() {
self.cancel.check()?;
let before = out.len();
if self.kind == JoinKind::Mark {
let marker = match &marks {
Some(marks) => marks[position].clone(),
None => self.marker(left_row, &left_types, scanned_over)?,
};
let mut row = pad_right(left_row, right_types.len());
let Some(position) = self.marker else {
return Err(Error::internal("a mark join has no marker column"));
};
row[left_types.len() + position] = marker;
out.push(row);
scratch.grow(out[before..].iter().map(|row| rows::footprint(row)).sum())?;
continue;
}
let scanned;
let hits: &[u32] = match (&index, &residual) {
(Some(index), Some(residual)) => {
index.matches(left_slots[position], &mut chain);
residual.keep(left_row, &right, &chain, &mut kept)?
}
_ => {
scanned = self.matching(left_row, &left_types, scanned_over)?;
&scanned
}
};
for &hit in hits {
matched[hit as usize] = true;
}
match self.kind {
JoinKind::Mark => unreachable!("mark joins leave before collecting hits"),
JoinKind::Semi => {
if !hits.is_empty() {
out.push(left_row.clone());
}
}
JoinKind::Anti => {
if hits.is_empty() {
out.push(left_row.clone());
}
}
JoinKind::Single => {
if hits.len() > 1 {
return Err(too_many_rows());
}
match hits.first() {
Some(&hit) => out.push(pair(left_row, &right.row(hit))),
None => out.push(pad_right(left_row, right_types.len())),
}
}
JoinKind::Left | JoinKind::Full if hits.is_empty() => {
out.push(pad_right(left_row, right_types.len()));
}
_ => {
for &hit in hits {
out.push(pair(left_row, &right.row(hit)));
}
}
}
scratch.grow(out[before..].iter().map(|row| rows::footprint(row)).sum())?;
}
if matches!(self.kind, JoinKind::Right | JoinKind::Full) {
for (at, seen) in matched.iter().enumerate() {
if !seen {
out.push(pad_left(left_types.len(), &right.row(at as u32)));
}
}
}
Ok(out)
}
fn matching(
&self,
left_row: &[Value],
left_types: &[LogicalType],
right_chunks: &[Chunk],
) -> Result<Vec<u32>> {
let mut hits = Vec::new();
let mut base: u32 = 0;
for chunk in right_chunks {
let rows = u32::try_from(chunk.len()).unwrap_or(PAD);
if self.conditions.is_empty() {
hits.extend(base..base + rows);
} else {
let combined = widen(left_row, left_types, chunk)?;
let flags = evaluate_all_in_time_zone(
self.plan,
&self.conditions,
&self.combined,
&combined,
self.time_zone,
)?;
let merged = combine(Connective::And, &flags)?;
for row in 0..rows {
if is_true(&merged.value_at(row as usize)) {
hits.push(base + row);
}
}
}
base += rows;
}
Ok(hits)
}
fn marker(
&self,
left_row: &[Value],
left_types: &[LogicalType],
right_chunks: &[Chunk],
) -> Result<Value> {
let mut unknown = false;
for chunk in right_chunks {
let combined = widen(left_row, left_types, chunk)?;
let flags = evaluate_all_in_time_zone(
self.plan,
&self.conditions,
&self.combined,
&combined,
self.time_zone,
)?;
let merged = combine(Connective::And, &flags)?;
for row in 0..chunk.len() {
match merged.value_at(row) {
Value::Boolean(true) => return Ok(Value::Boolean(true)),
Value::Null => unknown = true,
_ => {}
}
}
}
Ok(if unknown { Value::Null } else { Value::Boolean(false) })
}
fn marks(
&self,
left_rows: &[Vec<Value>],
right_chunks: &[Chunk],
scratch: &mut Reservation,
) -> Result<Option<Vec<Value>>> {
let Some(found) =
equalities(self.plan, &self.conditions, &self.left_schema, &self.right_schema)
else {
return Ok(None);
};
if found.left.len() != 1 || !found.residual.is_empty() {
return Ok(None);
}
if right_chunks.iter().all(Chunk::is_empty) {
return Ok(Some(vec![Value::Boolean(false); left_rows.len()]));
}
let nulls_are_values = found.null_is_a_value[0];
let gathered = found.gathered(self.plan, &self.right_schema, self.time_zone);
let undecided = !nulls_are_values && any_null_key(gathered, right_chunks, &self.cancel)?;
let index = lookup(gathered, right_chunks, &self.cancel, &Lease::alone(), scratch)?;
let types = self.left_schema.types();
scratch.grow(u64::try_from(left_rows.len()).unwrap_or(u64::MAX))?;
let mut marks = Vec::with_capacity(left_rows.len());
let mut probing = Scratch::default();
let mut slots = Vec::new();
for batch in left_rows.chunks(VECTOR_SIZE) {
self.cancel.check()?;
let chunk = rows::pack(&types, batch)?;
let columns = evaluate_all_in_time_zone(
self.plan,
&found.left,
&self.left_schema,
&chunk,
self.time_zone,
)?;
index.slots(&columns, batch.len(), &found.null_is_a_value, &mut probing, &mut slots);
for (row, &slot) in slots.iter().take(batch.len()).enumerate() {
marks.push(if slot != MISS {
Value::Boolean(true)
} else if undecided || (!nulls_are_values && columns[0].is_null_at(row)) {
Value::Null
} else {
Value::Boolean(false)
});
}
}
Ok(Some(marks))
}
}
impl Sink for Join<'_> {
type Local = Gathering;
fn local(&self) -> Gathering {
gather::gathering(&self.memory)
}
fn parallel(&self) -> bool {
false
}
fn sink(&self, chunk: &Chunk, local: &mut Gathering) -> Result<Progress> {
gather::take(chunk, local)?;
Ok(Progress::More)
}
fn combine(&self, local: Gathering) -> Result<()> {
let (rows, charged) = gather::into_parts(local);
self.left.lock().map_err(poisoned)?.extend(rows);
self.charged.lock().map_err(poisoned)?.push(charged);
Ok(())
}
fn finalize(&self, threads: &Lease<'_>) -> Result<()> {
let left_rows = std::mem::take(&mut *self.left.lock().map_err(poisoned)?);
let right_chunks = held(&self.right)?;
let mut out = self.joined(&left_rows, &right_chunks, threads)?;
if self.swapped {
let width = self.left_schema.width();
for row in &mut out {
row.rotate_left(width);
}
}
drop(left_rows);
drop(right_chunks);
let mut held = self.held.lock().map_err(poisoned)?;
let chunks = rows::chunks(&self.schema.types(), &out, &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 join gathered")
}
#[derive(Debug)]
pub(crate) struct CrossProduct {
left_types: Vec<LogicalType>,
types: Vec<LogicalType>,
schema: Schema,
right: Buffered,
}
#[derive(Debug)]
pub(crate) struct Crossing {
left: Option<Chunk>,
row: usize,
at: usize,
}
impl CrossProduct {
pub(crate) fn new(left: &Schema, right_schema: &Schema, right: Buffered) -> Self {
let schema = Schema::concat(left, right_schema);
Self { left_types: left.types(), types: schema.types(), schema, right }
}
pub(crate) fn schema(&self) -> &Schema {
&self.schema
}
}
impl Stream for CrossProduct {
type Local = Crossing;
fn local(&self) -> Crossing {
Crossing { left: None, row: 0, at: 0 }
}
fn push(&self, chunk: &mut Chunk, local: &mut Crossing) -> Result<Progress> {
let stored = self.right.len()?;
let left = match local.left.take() {
Some(left) => left,
None => {
local.row = 0;
local.at = 0;
chunk.clone()
}
};
if stored == 0 || left.is_empty() {
*chunk = Chunk::empty(&self.types);
return Ok(Progress::More);
}
let right = self
.right
.at(local.at)?
.ok_or_else(|| Error::internal("a cross product asked for a chunk nobody kept"))?;
let row: Vec<Value> = left.row(local.row).collect();
*chunk = widen(&row, &self.left_types, &right)?;
local.at += 1;
if local.at >= stored {
local.at = 0;
local.row += 1;
}
if local.row >= left.len() {
return Ok(Progress::More);
}
local.left = Some(left);
Ok(Progress::Again)
}
}
#[derive(Debug)]
pub(crate) struct Probe<'a> {
plan: &'a Plan,
kind: JoinKind,
equalities: Equalities,
marker: Option<usize>,
left_schema: Schema,
right_schema: Schema,
combined: Schema,
left_types: Vec<LogicalType>,
right_types: Vec<LogicalType>,
left_width: usize,
schema: Schema,
swapped: bool,
cancel: Cancel,
gathered: Buffered,
built: OnceLock<Result<Arc<Built>>>,
held: Mutex<Reservation>,
time_zone: SessionTimeZone,
}
#[derive(Debug)]
struct Built {
rows: Build,
index: Lookup,
undecided: bool,
}
const RESIDUAL_BATCH: usize = 16 * VECTOR_SIZE;
#[derive(Debug, Default)]
struct Candidates {
right: Vec<u32>,
start: Vec<u32>,
from: usize,
to: usize,
raw: Vec<u32>,
driving: Vec<u32>,
pass: Vec<bool>,
chain: Vec<u32>,
}
impl Candidates {
fn holds(&self, row: usize) -> bool {
row >= self.from && row < self.to
}
fn of(&self, row: usize) -> &[u32] {
let Some(index) = row.checked_sub(self.from) else {
return &[];
};
let (Some(&begin), Some(&end)) = (self.start.get(index), self.start.get(index + 1)) else {
return &[];
};
self.right.get(begin as usize..end as usize).unwrap_or_default()
}
fn forget(&mut self) {
self.from = 0;
self.to = 0;
}
fn fill(
&mut self,
residual: &Residual<'_>,
left: &Chunk,
built: &Built,
slots: &[usize],
) -> Result<()> {
self.raw.clear();
self.driving.clear();
self.start.clear();
self.right.clear();
self.start.push(0);
let mut row = self.from;
while row < left.len() {
let slot = slots.get(row).copied().unwrap_or(MISS);
built.index.matches(slot, &mut self.chain);
let at = u32::try_from(row).map_err(|_| too_many_rows())?;
self.raw.extend_from_slice(&self.chain);
self.driving.extend(std::iter::repeat_n(at, self.chain.len()));
row += 1;
self.start.push(u32::try_from(self.raw.len()).map_err(|_| too_many_rows())?);
if self.raw.len() >= RESIDUAL_BATCH {
break;
}
}
self.to = row;
residual.keeps(left, &built.rows, &self.driving, &self.raw, &mut self.pass)?;
let mut kept: usize = 0;
for index in 0..self.to - self.from {
let begin = self.start[index] as usize;
let end = self.start[index + 1] as usize;
self.start[index] = u32::try_from(kept).map_err(|_| too_many_rows())?;
for at in begin..end {
if self.pass.get(at).copied().unwrap_or(false) {
self.right.push(self.raw[at]);
kept += 1;
}
}
}
if let Some(last) = self.start.last_mut() {
*last = u32::try_from(kept).map_err(|_| too_many_rows())?;
}
Ok(())
}
}
#[derive(Debug)]
pub(crate) struct Probing {
left: Option<Chunk>,
keys: Vec<Vector>,
slots: Vec<usize>,
scratch: Scratch,
chain: Vec<u32>,
cand: Candidates,
left_at: Vec<u32>,
right_at: Vec<u32>,
row: usize,
hit: usize,
}
pub(crate) fn streamed(kind: JoinKind) -> bool {
matches!(
kind,
JoinKind::Inner
| JoinKind::Left
| JoinKind::Semi
| JoinKind::Anti
| JoinKind::Single
| JoinKind::Mark
)
}
impl<'a> Probe<'a> {
pub(crate) fn new(
plan: &'a Plan,
left: &Schema,
right: &Gathered<'_>,
kind: JoinKind,
conditions: Slice,
cancel: &Cancel,
memory: &Memory,
) -> Option<Self> {
if !streamed(kind) {
return None;
}
let swapped = right.swapped;
let right_schema = right.schema;
let equalities = equalities(plan, plan.expr_list(conditions), left, right_schema)?;
let marker = if kind == JoinKind::Mark {
if equalities.left.len() != 1 || !equalities.residual.is_empty() {
return None;
}
Some(right.marker.or_else(|| right_schema.bindings().len().checked_sub(1))?)
} else {
None
};
let schema = match kind {
JoinKind::Semi | JoinKind::Anti => left.clone(),
_ if swapped => Schema::concat(right_schema, left),
_ => Schema::concat(left, right_schema),
};
Some(Self {
plan,
kind,
equalities,
marker,
left_schema: left.clone(),
right_schema: right_schema.clone(),
combined: Schema::concat(left, right_schema),
left_types: left.types(),
right_types: right_schema.types(),
left_width: left.width(),
schema,
swapped,
cancel: cancel.clone(),
gathered: right.chunks.clone(),
built: OnceLock::new(),
held: Mutex::new(memory.reservation()),
time_zone: SessionTimeZone::default(),
})
}
#[must_use]
pub(crate) fn in_session(mut self, session: &Session) -> Self {
self.time_zone = session.session_time_zone();
self
}
pub(crate) fn schema(&self) -> &Schema {
&self.schema
}
pub(crate) fn sideways(&self) -> Option<(ExprRef, ColumnBinding)> {
if !matches!(self.kind, JoinKind::Inner | JoinKind::Semi) {
return None;
}
self.keyed_sideways()
}
fn keyed_sideways(&self) -> Option<(ExprRef, ColumnBinding)> {
let at = self.equalities.null_is_a_value.iter().position(|&stored| !stored)?;
let Expr::Column(binding) = *self.plan.expr(*self.equalities.left.get(at)?) else {
return None;
};
Some((*self.equalities.right.get(at)?, binding))
}
fn marked(
&self,
chunk: &mut Chunk,
left: &Chunk,
built: &Built,
local: &Probing,
) -> Result<Progress> {
let rows = left.len();
let empty = built.rows.rows() == 0;
let driving = match self.equalities.null_is_a_value.first() {
Some(true) => None,
_ => local.keys.first(),
};
let mut marks = vec![false; rows];
let mut known = vec![true; rows];
for (row, (mark, decided)) in marks.iter_mut().zip(known.iter_mut()).enumerate() {
if local.slots.get(row).copied().unwrap_or(MISS) != MISS {
*mark = true;
} else if !empty {
*decided = !built.undecided && !driving.is_some_and(|key| key.is_null_at(row));
}
}
let marker = Vector::flat(LogicalType::Boolean, Data::Bool(marks.into()))?
.with_validity(Validity::from_run(&known));
let mut columns: Vec<Vector> = left.columns().to_vec();
for logical in &self.right_types {
columns.push(Vector::constant(logical.clone(), Value::Null, rows));
}
let at = self.marker.map_or(usize::MAX, |at| self.left_width + at);
let Some(slot) = columns.get_mut(at) else {
return Err(Error::internal("a mark join has no marker column"));
};
*slot = marker;
*chunk = Chunk::with_rows(columns, rows)?;
Ok(Progress::More)
}
fn residual(&self) -> Residual<'_> {
Residual {
plan: self.plan,
exprs: &self.equalities.residual,
combined: &self.combined,
left_types: &self.left_types,
time_zone: self.time_zone,
}
}
fn built(&self) -> Result<Arc<Built>> {
self.built_with(&Lease::alone())
}
fn built_with(&self, threads: &Lease<'_>) -> Result<Arc<Built>> {
self.built
.get_or_init(|| {
let chunks = held(&self.gathered)?;
let keying =
self.equalities.gathered(self.plan, &self.right_schema, self.time_zone);
let mut charged = self.held.lock().map_err(poisoned)?;
let undecided = self.kind == JoinKind::Mark
&& !self.equalities.null_is_a_value.first().copied().unwrap_or(false)
&& any_null_key(keying, &chunks, &self.cancel)?;
let index = lookup(keying, &chunks, &self.cancel, threads, &mut charged)?;
let rows = Build::new(&self.right_types, &chunks, threads)?;
charged.grow(rows.footprint())?;
Ok(Arc::new(Built { rows, index, undecided }))
})
.clone()
}
}
impl Stream for Probe<'_> {
type Local = Probing;
fn local(&self) -> Probing {
Probing {
left: None,
keys: Vec::new(),
slots: Vec::new(),
scratch: Scratch::default(),
chain: Vec::new(),
cand: Candidates::default(),
left_at: Vec::new(),
right_at: Vec::new(),
row: 0,
hit: 0,
}
}
fn prepare(&self, threads: &Lease<'_>) -> Result<()> {
self.built_with(threads)?;
Ok(())
}
fn push(&self, chunk: &mut Chunk, local: &mut Probing) -> Result<Progress> {
let built = self.built()?;
let left = match local.left.take() {
Some(left) => left,
None => {
local.row = 0;
local.hit = 0;
local.cand.forget();
let left = chunk.clone();
local.keys = if built.index.is_empty() {
Vec::new()
} else {
evaluate_all_in_time_zone(
self.plan,
&self.equalities.left,
&self.left_schema,
&left,
self.time_zone,
)?
};
local.slots.clear();
if !local.keys.is_empty() {
built.index.slots(
&local.keys,
left.len(),
&self.equalities.null_is_a_value,
&mut local.scratch,
&mut local.slots,
);
}
left
}
};
if self.kind == JoinKind::Mark {
return self.marked(chunk, &left, &built, local);
}
let residual = self.residual();
{
let Probing { slots, chain, cand, left_at, right_at, row, hit, .. } = &mut *local;
left_at.clear();
right_at.clear();
while *row < left.len() && left_at.len() < VECTOR_SIZE {
self.cancel.check()?;
let found: &[u32] = if residual.exprs.is_empty() {
let slot = slots.get(*row).copied().unwrap_or(MISS);
built.index.matches(slot, chain);
chain
} else {
if !cand.holds(*row) {
cand.from = *row;
cand.fill(&residual, &left, &built, slots)?;
}
cand.of(*row)
};
let at = u32::try_from(*row).map_err(|_| too_many_rows())?;
match self.kind {
JoinKind::Semi => {
if !found.is_empty() {
left_at.push(at);
}
}
JoinKind::Anti => {
if found.is_empty() {
left_at.push(at);
}
}
JoinKind::Single => {
if found.len() > 1 {
return Err(too_many_rows());
}
left_at.push(at);
right_at.push(found.first().copied().unwrap_or(PAD));
}
JoinKind::Left if found.is_empty() => {
left_at.push(at);
right_at.push(PAD);
}
_ => {
let room = VECTOR_SIZE - left_at.len();
let end = (*hit + room).min(found.len());
for &found_at in &found[*hit..end] {
left_at.push(at);
right_at.push(found_at);
}
if end < found.len() {
*hit = end;
break;
}
*hit = 0;
}
}
*row += 1;
}
}
let mut columns: Vec<Vector> = left
.columns()
.iter()
.map(|column| column.gather(&local.left_at))
.collect::<Result<Vec<_>>>()?;
if !matches!(self.kind, JoinKind::Semi | JoinKind::Anti) {
columns.extend(built.rows.gather(&local.right_at)?);
}
if self.swapped {
columns.rotate_left(self.left_width);
}
*chunk = Chunk::with_rows(columns, local.left_at.len())?;
if local.row < left.len() {
local.left = Some(left);
return Ok(Progress::Again);
}
Ok(Progress::More)
}
}
#[derive(Debug)]
pub(crate) struct Marking<'a> {
probe: Probe<'a>,
schema: Schema,
marked: Mutex<Vec<u64>>,
held: Mutex<Reservation>,
out: Buffered,
}
#[derive(Debug)]
pub(crate) struct Marks {
keys: Vec<Vector>,
slots: Vec<usize>,
scratch: Scratch,
chain: Vec<u32>,
cand: Candidates,
bits: Vec<u64>,
}
impl<'a> Marking<'a> {
pub(crate) fn new(
plan: &'a Plan,
left: &Schema,
right: &Gathered<'_>,
kind: JoinKind,
conditions: Slice,
cancel: &Cancel,
memory: &Memory,
) -> Option<(Self, Buffered)> {
if !matches!(kind, JoinKind::Semi | JoinKind::Anti) || !right.swapped {
return None;
}
let probe = Probe::new(plan, left, right, kind, conditions, cancel, memory)?;
let out = Buffered::new();
let marking = Self {
probe,
schema: right.schema.clone(),
marked: Mutex::new(Vec::new()),
held: Mutex::new(memory.reservation()),
out: out.clone(),
};
Some((marking, out))
}
#[must_use]
pub(crate) fn in_session(mut self, session: &Session) -> Self {
self.probe = self.probe.in_session(session);
self
}
pub(crate) fn schema(&self) -> &Schema {
&self.schema
}
pub(crate) fn sideways(&self) -> Option<(ExprRef, ColumnBinding)> {
self.probe.keyed_sideways()
}
}
impl Sink for Marking<'_> {
type Local = Marks;
fn local(&self) -> Marks {
Marks {
keys: Vec::new(),
slots: Vec::new(),
scratch: Scratch::default(),
chain: Vec::new(),
cand: Candidates::default(),
bits: Vec::new(),
}
}
fn prepare(&self, threads: &Lease<'_>) -> Result<()> {
self.probe.built_with(threads)?;
Ok(())
}
fn sink(&self, chunk: &Chunk, local: &mut Marks) -> Result<Progress> {
let built = self.probe.built()?;
let rows = built.rows.rows();
if rows == 0 || chunk.is_empty() {
return Ok(Progress::More);
}
if local.bits.is_empty() {
local.bits = vec![0; rows.div_ceil(u64::BITS as usize)];
}
local.keys = if built.index.is_empty() {
Vec::new()
} else {
evaluate_all_in_time_zone(
self.probe.plan,
&self.probe.equalities.left,
&self.probe.left_schema,
chunk,
self.probe.time_zone,
)?
};
local.slots.clear();
if !local.keys.is_empty() {
built.index.slots(
&local.keys,
chunk.len(),
&self.probe.equalities.null_is_a_value,
&mut local.scratch,
&mut local.slots,
);
}
local.cand.forget();
let residual = self.probe.residual();
let Marks { slots, chain, cand, bits, .. } = local;
for row in 0..chunk.len() {
let found: &[u32] = if residual.exprs.is_empty() {
let slot = slots.get(row).copied().unwrap_or(MISS);
built.index.matches(slot, chain);
chain
} else {
if !cand.holds(row) {
cand.from = row;
cand.fill(&residual, chunk, &built, slots)?;
}
cand.of(row)
};
for &at in found {
let at = at as usize;
if let Some(word) = bits.get_mut(at / u64::BITS as usize) {
*word |= 1 << (at % u64::BITS as usize);
}
}
}
Ok(Progress::More)
}
fn combine(&self, local: Marks) -> Result<()> {
if local.bits.is_empty() {
return Ok(());
}
let mut marked = self.marked.lock().map_err(poisoned)?;
if marked.is_empty() {
*marked = local.bits;
return Ok(());
}
for (word, one) in marked.iter_mut().zip(local.bits) {
*word |= one;
}
Ok(())
}
fn finalize(&self, threads: &Lease<'_>) -> Result<()> {
let built = self.probe.built_with(threads)?;
let marked = std::mem::take(&mut *self.marked.lock().map_err(poisoned)?);
let wanted = self.probe.kind == JoinKind::Semi;
let mut chunks = Vec::new();
let mut at: Vec<u32> = Vec::with_capacity(VECTOR_SIZE);
let mut charged = self.held.lock().map_err(poisoned)?;
for row in 0..built.rows.rows() {
let bit = marked
.get(row / u64::BITS as usize)
.is_some_and(|word| word >> (row % u64::BITS as usize) & 1 == 1);
if bit != wanted {
continue;
}
at.push(u32::try_from(row).map_err(|_| unaddressable())?);
if at.len() == VECTOR_SIZE {
let chunk = built.rows.chunk(&at)?;
charged.grow(chunk.footprint() as u64)?;
chunks.push(chunk);
at.clear();
}
}
if !at.is_empty() {
let chunk = built.rows.chunk(&at)?;
charged.grow(chunk.footprint() as u64)?;
chunks.push(chunk);
}
self.out.fill(chunks)
}
}
#[derive(Debug)]
pub(crate) struct Padding<'a> {
probe: Probe<'a>,
marked: OnceLock<Vec<AtomicU64>>,
}
#[derive(Debug)]
pub(crate) struct Padded {
probing: Probing,
}
impl<'a> Padding<'a> {
pub(crate) fn new(
plan: &'a Plan,
left: &Schema,
right: &Gathered<'_>,
kind: JoinKind,
conditions: Slice,
cancel: &Cancel,
memory: &Memory,
) -> Option<Self> {
let pairing = match kind {
JoinKind::Right => JoinKind::Inner,
JoinKind::Full => JoinKind::Left,
_ => return None,
};
let probe = Probe::new(plan, left, right, pairing, conditions, cancel, memory)?;
Some(Self { probe, marked: OnceLock::new() })
}
#[must_use]
pub(crate) fn in_session(mut self, session: &Session) -> Self {
self.probe = self.probe.in_session(session);
self
}
pub(crate) fn schema(&self) -> &Schema {
self.probe.schema()
}
pub(crate) fn sideways(&self) -> Option<(ExprRef, ColumnBinding)> {
None
}
fn mark(&self, at: &[u32]) {
let Some(marked) = self.marked.get() else {
return;
};
for &row in at {
if row == PAD {
continue;
}
let row = row as usize;
if let Some(word) = marked.get(row / u64::BITS as usize) {
word.fetch_or(1 << (row % u64::BITS as usize), Ordering::Relaxed);
}
}
}
fn padded(&self, built: &Built, at: &[u32]) -> Result<Chunk> {
let mut columns: Vec<Vector> = self
.probe
.left_types
.iter()
.map(|ty| Vector::constant(ty.clone(), Value::Null, at.len()))
.collect();
columns.extend(built.rows.gather(at)?);
if self.probe.swapped {
columns.rotate_left(self.probe.left_width);
}
Chunk::with_rows(columns, at.len())
}
}
impl Stream for Padding<'_> {
type Local = Padded;
fn local(&self) -> Padded {
Padded { probing: Stream::local(&self.probe) }
}
fn prepare(&self, threads: &Lease<'_>) -> Result<()> {
let built = self.probe.built_with(threads)?;
let words = built.rows.rows().div_ceil(u64::BITS as usize);
let _ = self.marked.set((0..words).map(|_| AtomicU64::new(0)).collect());
Ok(())
}
fn drains(&self) -> bool {
true
}
fn drain(&self, out: &mut dyn FnMut(&mut Chunk) -> Result<Progress>) -> Result<()> {
let built = self.probe.built()?;
let rows = built.rows.rows();
let empty = Vec::new();
let marked = self.marked.get().unwrap_or(&empty);
let mut at: Vec<u32> = Vec::with_capacity(VECTOR_SIZE);
for (word, bits) in marked.iter().enumerate() {
let bits = bits.load(Ordering::Relaxed);
if bits == u64::MAX {
continue;
}
let first = word * u64::BITS as usize;
for bit in 0..u64::BITS as usize {
let row = first + bit;
if row >= rows {
break;
}
if bits >> bit & 1 == 1 {
continue;
}
at.push(u32::try_from(row).map_err(|_| unaddressable())?);
if at.len() == VECTOR_SIZE {
let mut chunk = self.padded(&built, &at)?;
at.clear();
if out(&mut chunk)? == Progress::Done {
return Ok(());
}
}
}
}
if !at.is_empty() {
let mut chunk = self.padded(&built, &at)?;
out(&mut chunk)?;
}
Ok(())
}
fn push(&self, chunk: &mut Chunk, local: &mut Padded) -> Result<Progress> {
let progress = self.probe.push(chunk, &mut local.probing)?;
self.mark(&local.probing.right_at);
Ok(progress)
}
}
fn unaddressable() -> Error {
Error::internal("a join gathered more rows than it can address")
}
fn too_many_rows() -> Error {
Error::invalid_input(
"More than one row returned by a subquery used as an expression - scalar subqueries can only return a single row.\n\nUse \"SET scalar_subquery_error_on_multiple_rows=false\" to revert to previous behavior of returning a random row."
.to_string(),
)
}
#[derive(Debug)]
struct Equalities {
left: Vec<ExprRef>,
right: Vec<ExprRef>,
null_is_a_value: Vec<bool>,
residual: Vec<ExprRef>,
}
impl Equalities {
fn driving<'a>(
&'a self,
plan: &'a Plan,
schema: &'a Schema,
time_zone: SessionTimeZone,
) -> Keying<'a> {
Keying { plan, exprs: &self.left, schema, nulls: &self.null_is_a_value, time_zone }
}
fn gathered<'a>(
&'a self,
plan: &'a Plan,
schema: &'a Schema,
time_zone: SessionTimeZone,
) -> Keying<'a> {
Keying { plan, exprs: &self.right, schema, nulls: &self.null_is_a_value, time_zone }
}
}
#[derive(Debug, Clone, Copy)]
struct Keying<'a> {
plan: &'a Plan,
exprs: &'a [ExprRef],
schema: &'a Schema,
nulls: &'a [bool],
time_zone: SessionTimeZone,
}
#[derive(Debug, Clone, Copy)]
struct Residual<'a> {
plan: &'a Plan,
exprs: &'a [ExprRef],
combined: &'a Schema,
left_types: &'a [LogicalType],
time_zone: SessionTimeZone,
}
impl Residual<'_> {
fn keep<'h>(
&self,
left_row: &[Value],
rows: &Build,
hits: &'h [u32],
into: &'h mut Vec<u32>,
) -> Result<&'h [u32]> {
if self.exprs.is_empty() || hits.is_empty() {
return Ok(hits);
}
into.clear();
for batch in hits.chunks(VECTOR_SIZE) {
let chunk = rows.chunk(batch)?;
let combined = widen(left_row, self.left_types, &chunk)?;
let flags = evaluate_all_in_time_zone(
self.plan,
self.exprs,
self.combined,
&combined,
self.time_zone,
)?;
let merged = combine(Connective::And, &flags)?;
for (at, &hit) in batch.iter().enumerate() {
if is_true(&merged.value_at(at)) {
into.push(hit);
}
}
}
Ok(into)
}
fn keeps(
&self,
left: &Chunk,
rows: &Build,
driving: &[u32],
gathered: &[u32],
into: &mut Vec<bool>,
) -> Result<()> {
into.clear();
if self.exprs.is_empty() {
into.resize(gathered.len(), true);
return Ok(());
}
into.reserve(gathered.len());
let mut at = 0;
while at < gathered.len() {
let end = (at + VECTOR_SIZE).min(gathered.len());
let mut columns: Vec<Vector> = left
.columns()
.iter()
.map(|column| column.gather(&driving[at..end]))
.collect::<Result<Vec<_>>>()?;
columns.extend(rows.gather(&gathered[at..end])?);
let combined = Chunk::with_rows(columns, end - at)?;
let flags = evaluate_all_in_time_zone(
self.plan,
self.exprs,
self.combined,
&combined,
self.time_zone,
)?;
let merged = combine(Connective::And, &flags)?;
for row in 0..end - at {
into.push(is_true(&merged.value_at(row)));
}
at = end;
}
Ok(())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Side {
Driving,
Gathered,
}
fn side_of(plan: &Plan, expr: ExprRef, driving: &Schema, gathered: &Schema) -> Option<Side> {
let mut side = None;
let mut mixed = false;
columns(plan, expr, &mut |binding| {
let found = if driving.position_of(binding).is_some() {
Some(Side::Driving)
} else if gathered.position_of(binding).is_some() {
Some(Side::Gathered)
} else {
None
};
match (side, found) {
(_, None) => mixed = true,
(None, Some(one)) => side = Some(one),
(Some(held), Some(one)) => mixed |= held != one,
}
});
if mixed { None } else { side }
}
fn columns(plan: &Plan, expr: ExprRef, found: &mut impl FnMut(ColumnBinding)) {
match *plan.expr(expr) {
Expr::Column(binding) => found(binding),
Expr::Constant(_) => {}
Expr::Cast { input, .. } => columns(plan, input, found),
Expr::Compare { left, right, .. } => {
columns(plan, left, found);
columns(plan, right, found);
}
Expr::Conjunction { children, .. } | Expr::Function { args: children, .. } => {
for &child in plan.expr_list(children) {
columns(plan, child, found);
}
}
Expr::Aggregate { args, filter, .. } | Expr::Window { args, filter, .. } => {
for &arg in plan.expr_list(args) {
columns(plan, arg, found);
}
if let Some(inner) = filter {
columns(plan, inner, found);
}
}
Expr::Case { arms, otherwise } => {
for arm in plan.arm_list(arms) {
columns(plan, arm.when, found);
columns(plan, arm.then, found);
}
if let Some(inner) = otherwise {
columns(plan, inner, found);
}
}
}
}
fn equalities(
plan: &Plan,
conditions: &[ExprRef],
left_schema: &Schema,
right_schema: &Schema,
) -> Option<Equalities> {
let mut found = Equalities {
left: Vec::new(),
right: Vec::new(),
null_is_a_value: Vec::new(),
residual: Vec::new(),
};
for &condition in conditions {
let Expr::Compare { op: op @ (CompareOp::Equal | CompareOp::NotDistinctFrom), left, right } =
*plan.expr(condition)
else {
found.residual.push(condition);
continue;
};
if plan.expr_type(left) != plan.expr_type(right) || !plan.expr_type(left).is_keyed() {
found.residual.push(condition);
continue;
}
match (
side_of(plan, left, left_schema, right_schema),
side_of(plan, right, left_schema, right_schema),
) {
(Some(Side::Driving), Some(Side::Gathered)) => {
found.left.push(left);
found.right.push(right);
}
(Some(Side::Gathered), Some(Side::Driving)) => {
found.left.push(right);
found.right.push(left);
}
_ => {
found.residual.push(condition);
continue;
}
}
found.null_is_a_value.push(op == CompareOp::NotDistinctFrom);
}
(!found.left.is_empty()).then_some(found)
}
fn lookup(
keying: Keying<'_>,
chunks: &[Chunk],
cancel: &Cancel,
threads: &Lease<'_>,
scratch: &mut Reservation,
) -> Result<Lookup> {
let Keying { plan, exprs, schema, nulls, time_zone } = keying;
let rows: usize = chunks.iter().map(Chunk::len).sum();
let mut keyed: Vec<Chunk> = Vec::with_capacity(chunks.len());
for chunk in chunks {
cancel.check()?;
let columns = evaluate_all_in_time_zone(plan, exprs, schema, chunk, time_zone)?;
keyed.push(Chunk::with_rows(columns, chunk.len())?);
}
let Some(types) = keyed.first().map(Chunk::types) else {
return Lookup::build(&[], 0, nulls, threads, cancel);
};
let keys = laid_out(&types, &keyed, threads)?;
drop(keyed);
let lookup = Lookup::build(&keys, rows, nulls, threads, cancel)?;
scratch.grow(lookup.footprint())?;
Ok(lookup)
}
fn any_null_key(keying: Keying<'_>, chunks: &[Chunk], cancel: &Cancel) -> Result<bool> {
let Keying { plan, exprs, schema, time_zone, .. } = keying;
for chunk in chunks {
cancel.check()?;
let columns = evaluate_all_in_time_zone(plan, exprs, schema, chunk, time_zone)?;
if columns.iter().any(|column| column.validity().has_nulls(chunk.len())) {
return Ok(true);
}
}
Ok(false)
}
fn held(chunks: &Buffered) -> Result<Vec<Chunk>> {
let reader = chunks.reader();
(0..reader.len()?)
.map(|at| {
reader.at(at)?.ok_or_else(|| {
Error::internal("a join was given fewer gathered chunks than it was told about")
})
})
.collect()
}
fn found(
keying: Keying<'_>,
index: &Lookup,
rows: &[Vec<Value>],
cancel: &Cancel,
scratch: &mut Reservation,
) -> Result<Vec<usize>> {
let Keying { plan, exprs, schema, nulls, time_zone } = keying;
let types = schema.types();
let mut built = Vec::with_capacity(rows.len());
scratch
.grow(u64::try_from(rows.len().saturating_mul(size_of::<usize>())).unwrap_or(u64::MAX))?;
let mut probing = Scratch::default();
let mut slots = Vec::new();
for batch in rows.chunks(VECTOR_SIZE) {
cancel.check()?;
let chunk = rows::pack(&types, batch)?;
let columns = evaluate_all_in_time_zone(plan, exprs, schema, &chunk, time_zone)?;
index.slots(&columns, batch.len(), nulls, &mut probing, &mut slots);
built.extend_from_slice(&slots);
}
Ok(built)
}
fn pair(left: &[Value], right: &[Value]) -> Vec<Value> {
let mut row = left.to_vec();
row.extend(right.iter().cloned());
row
}
fn pad_right(left: &[Value], width: usize) -> Vec<Value> {
let mut row = left.to_vec();
row.extend(std::iter::repeat_n(Value::Null, width));
row
}
fn pad_left(width: usize, right: &[Value]) -> Vec<Value> {
let mut row = vec![Value::Null; width];
row.extend(right.iter().cloned());
row
}
fn positional(
left: &[Vec<Value>],
right: &[Vec<Value>],
left_width: usize,
right_width: usize,
) -> Vec<Vec<Value>> {
let rows = left.len().max(right.len());
(0..rows)
.map(|at| match (left.get(at), right.get(at)) {
(Some(left), Some(right)) => pair(left, right),
(Some(left), None) => pad_right(left, right_width),
(None, Some(right)) => pad_left(left_width, right),
(None, None) => Vec::new(),
})
.collect()
}
fn widen(left_row: &[Value], left_types: &[LogicalType], right: &Chunk) -> Result<Chunk> {
let rows = right.len();
let mut columns: Vec<Vector> = left_row
.iter()
.zip(left_types)
.map(|(value, ty)| Vector::constant(ty.clone(), value.clone(), rows))
.collect();
columns.extend(right.columns().iter().cloned());
Chunk::with_rows(columns, rows)
}
#[cfg(test)]
mod tests {
use rudb_common::{Cancel, Field, LogicalType, Memory, Value};
use rudb_plan::{ColumnBinding, CompareOp, Expr, ExprRef, JoinKind, Plan, Slice};
use rudb_vector::{Data, Validity, Vector};
use super::{
Buffered, Chunk, CrossProduct, Gathered, Join, Marking, Padding, Probe, Progress, Schema,
Side, Sink, Stream, VECTOR_SIZE, equalities, side_of,
};
use crate::gather::Keep;
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 wide_chunk(values: &[i64]) -> Chunk {
let column = Vector::flat(LogicalType::BigInt, Data::Int64(values.to_vec().into()))
.expect("big integers are an i64 layout");
Chunk::new(vec![column]).expect("one column is one length")
}
fn schema(name: &str, table: u32) -> Schema {
Schema::numbered(vec![Field::new(name, LogicalType::Integer)], table)
}
fn typed_schema(name: &str, table: u32, ty: LogicalType) -> Schema {
Schema::numbered(vec![Field::new(name, ty)], table)
}
fn column(plan: &mut Plan, table: u32, ty: LogicalType) -> ExprRef {
plan.add_expr(Expr::Column(ColumnBinding::new(table, 0)), ty)
}
fn equal(plan: &mut Plan, left: ExprRef, right: ExprRef) -> ExprRef {
plan.add_expr(Expr::Compare { op: CompareOp::Equal, left, right }, LogicalType::Boolean)
}
fn greater(plan: &mut Plan, left: ExprRef, right: ExprRef) -> ExprRef {
plan.add_expr(Expr::Compare { op: CompareOp::Greater, left, right }, LogicalType::Boolean)
}
fn pair_schema(table: u32) -> Schema {
Schema::numbered(
vec![Field::new("k", LogicalType::Integer), Field::new("g", LogicalType::Integer)],
table,
)
}
fn column_at(plan: &mut Plan, table: u32, position: u32, ty: LogicalType) -> ExprRef {
plan.add_expr(Expr::Column(ColumnBinding::new(table, position)), ty)
}
fn pair_chunk(rows: &[(i32, i32)]) -> Chunk {
let each = |pick: fn(&(i32, i32)) -> i32| {
Vector::flat(
LogicalType::Integer,
Data::Int32(rows.iter().map(pick).collect::<Vec<i32>>().into()),
)
.expect("integers are an i32 layout")
};
Chunk::new(vec![each(|row| row.0), each(|row| row.1)]).expect("two columns of one length")
}
fn sides() -> (Schema, Schema) {
(typed_schema("a", 0, LogicalType::Integer), typed_schema("b", 1, LogicalType::BigInt))
}
fn probed(probe: &Probe<'_>, driving: &Chunk, width: usize) -> Vec<Vec<Value>> {
let mut local = probe.local();
let mut chunk = driving.clone();
let mut out = Vec::new();
loop {
let progress = probe.push(&mut chunk, &mut local).expect("a chunk");
out.extend((0..chunk.len()).map(|row| {
(0..width).map(|column| chunk.value_at(row, column)).collect::<Vec<Value>>()
}));
if progress != Progress::Again {
return out;
}
chunk = Chunk::empty(&[]);
}
}
fn some_column(values: &[Option<i32>]) -> Vector {
let held: Vec<i32> = values.iter().map(|value| value.unwrap_or_default()).collect();
let valid: Vec<bool> = values.iter().map(Option::is_some).collect();
Vector::flat(LogicalType::Integer, Data::Int32(held.into()))
.expect("integers are an i32 layout")
.with_validity(Validity::from_run(&valid))
}
fn some_chunk(values: &[Option<i32>]) -> Chunk {
Chunk::with_rows(vec![some_column(values)], values.len()).expect("one column is one length")
}
fn marked_schema(table: u32) -> Schema {
Schema::numbered(
vec![Field::new("k", LogicalType::Integer), Field::new("mark", LogicalType::Boolean)],
table,
)
}
fn marked_chunk(keys: &[Option<i32>]) -> Chunk {
let mark = Vector::constant(LogicalType::Boolean, Value::Null, keys.len());
Chunk::with_rows(vec![some_column(keys), mark], keys.len())
.expect("two columns of one length")
}
fn markers(gathered: &[Option<i32>], driving: &[Option<i32>]) -> Vec<Value> {
let mut plan = Plan::new();
let (left, right) = (schema("a", 0), marked_schema(1));
let conditions = {
let one = column(&mut plan, 0, LogicalType::Integer);
let other = column_at(&mut plan, 1, 0, LogicalType::Integer);
let key = equal(&mut plan, one, other);
plan.add_expr_list(&[key])
};
let memory = Memory::unlimited();
let (keep, rows) = Keep::new(&memory);
let mut local = keep.local();
if !gathered.is_empty() {
keep.sink(&marked_chunk(gathered), &mut local).expect("the gathered rows");
}
keep.combine(local).expect("the one instance");
keep.finalize(&rudb_pipeline::Lease::alone()).expect("the chunks");
let probe = Probe::new(
&plan,
&left,
&Gathered { schema: &right, chunks: rows, marker: Some(1), swapped: false },
JoinKind::Mark,
conditions,
&Cancel::new(),
&memory,
)
.expect("one equality is enough to mark on");
probed(&probe, &some_chunk(driving), 3).into_iter().map(|row| row[2].clone()).collect()
}
#[test]
fn a_mark_join_over_a_gathered_side_with_a_null_key_marks_every_miss_null() {
assert_eq!(
markers(&[Some(2), None], &[Some(2), Some(3), None]),
[Value::Boolean(true), Value::Null, Value::Null]
);
}
#[test]
fn a_mark_join_over_a_side_with_no_null_key_marks_a_miss_false() {
assert_eq!(
markers(&[Some(2)], &[Some(2), Some(3), None]),
[Value::Boolean(true), Value::Boolean(false), Value::Null]
);
}
#[test]
fn a_mark_join_over_an_empty_gathered_side_marks_everything_false() {
assert_eq!(markers(&[], &[Some(2), None]), [Value::Boolean(false), Value::Boolean(false)]);
}
#[test]
fn a_mark_join_over_a_side_of_nothing_but_nulls_marks_everything_null() {
assert_eq!(markers(&[None, None], &[Some(2), None]), [Value::Null, Value::Null]);
}
#[test]
fn a_mark_join_on_two_equalities_is_not_streamed() {
let mut plan = Plan::new();
let (left, right) = (pair_schema(0), pair_schema(1));
let conditions = {
let one = column_at(&mut plan, 0, 0, LogicalType::Integer);
let other = column_at(&mut plan, 1, 0, LogicalType::Integer);
let first = equal(&mut plan, one, other);
let above = column_at(&mut plan, 0, 1, LogicalType::Integer);
let below = column_at(&mut plan, 1, 1, LogicalType::Integer);
let second = equal(&mut plan, above, below);
plan.add_expr_list(&[first, second])
};
let memory = Memory::unlimited();
let (_keep, rows) = gathered(&memory, &[]);
assert!(
Probe::new(
&plan,
&left,
&Gathered { schema: &right, chunks: rows, marker: Some(1), swapped: false },
JoinKind::Mark,
conditions,
&Cancel::new(),
&memory,
)
.is_none()
);
}
fn gathered(memory: &Memory, values: &[i32]) -> (Keep<'static>, Buffered) {
let (keep, chunks) = Keep::new(memory);
let mut local = keep.local();
if !values.is_empty() {
keep.sink(&chunk(values), &mut local).expect("the right rows");
}
keep.combine(local).expect("the one instance");
keep.finalize(&rudb_pipeline::Lease::alone()).expect("the chunks");
(keep, chunks)
}
fn run(join: &Join<'_>, left: &[i32]) {
let mut local = join.local();
if !left.is_empty() {
join.sink(&chunk(left), &mut local).expect("the left rows");
}
join.combine(local).expect("the one instance");
join.finalize(&rudb_pipeline::Lease::alone()).expect("the answer");
}
fn rows(out: &Buffered, width: usize) -> Vec<Vec<Value>> {
let Some(chunk) = out.at(0).expect("readable") else { return Vec::new() };
(0..chunk.len())
.map(|row| (0..width).map(|column| chunk.value_at(row, column)).collect())
.collect()
}
#[test]
fn every_left_row_meets_every_right_row_when_there_is_no_condition() {
let plan = Plan::new();
let memory = Memory::unlimited();
let (_gather, right) = gathered(&memory, &[10, 20]);
let (join, out) = Join::new(
&plan,
&schema("a", 0),
Gathered { schema: &schema("b", 1), chunks: right, marker: None, swapped: false },
JoinKind::Inner,
Slice::EMPTY,
&Cancel::new(),
&memory,
);
run(&join, &[1, 2]);
assert_eq!(
rows(&out, 2),
[
vec![Value::Integer(1), Value::Integer(10)],
vec![Value::Integer(1), Value::Integer(20)],
vec![Value::Integer(2), Value::Integer(10)],
vec![Value::Integer(2), Value::Integer(20)],
]
);
}
#[test]
fn a_swapped_join_produces_the_plans_columns_in_the_plans_order() {
let plan = Plan::new();
let memory = Memory::unlimited();
let (_gather, gathered_side) = gathered(&memory, &[1, 2]);
let (join, out) = Join::new(
&plan,
&schema("b", 1),
Gathered {
schema: &schema("a", 0),
chunks: gathered_side,
marker: None,
swapped: true,
},
JoinKind::Inner,
Slice::EMPTY,
&Cancel::new(),
&memory,
);
run(&join, &[10, 20]);
assert_eq!(join.schema().position_of(ColumnBinding::new(0, 0)), Some(0));
assert_eq!(join.schema().position_of(ColumnBinding::new(1, 0)), Some(1));
assert_eq!(
rows(&out, 2),
[
vec![Value::Integer(1), Value::Integer(10)],
vec![Value::Integer(2), Value::Integer(10)],
vec![Value::Integer(1), Value::Integer(20)],
vec![Value::Integer(2), Value::Integer(20)],
]
);
}
#[test]
fn a_swapped_outer_join_pads_the_side_the_plan_called_the_right_one() {
let plan = Plan::new();
let memory = Memory::unlimited();
let (_gather, gathered_side) = gathered(&memory, &[1, 2]);
let (join, out) = Join::new(
&plan,
&schema("b", 1),
Gathered {
schema: &schema("a", 0),
chunks: gathered_side,
marker: None,
swapped: true,
},
JoinKind::Right,
Slice::EMPTY,
&Cancel::new(),
&memory,
);
run(&join, &[]);
assert_eq!(
rows(&out, 2),
[vec![Value::Integer(1), Value::Null], vec![Value::Integer(2), Value::Null],]
);
}
#[test]
fn an_anti_join_against_nothing_keeps_every_left_row() {
let plan = Plan::new();
let memory = Memory::unlimited();
let (_gather, right) = gathered(&memory, &[]);
let (join, out) = Join::new(
&plan,
&schema("a", 0),
Gathered { schema: &schema("b", 1), chunks: right, marker: None, swapped: false },
JoinKind::Anti,
Slice::EMPTY,
&Cancel::new(),
&memory,
);
run(&join, &[1, 2]);
assert_eq!(rows(&out, 1), [vec![Value::Integer(1)], vec![Value::Integer(2)]]);
}
#[test]
fn a_positional_join_pads_the_shorter_side() {
let plan = Plan::new();
let memory = Memory::unlimited();
let (_gather, right) = gathered(&memory, &[10]);
let (join, out) = Join::new(
&plan,
&schema("a", 0),
Gathered { schema: &schema("b", 1), chunks: right, marker: None, swapped: false },
JoinKind::Positional,
Slice::EMPTY,
&Cancel::new(),
&memory,
);
run(&join, &[1, 2]);
assert_eq!(
rows(&out, 2),
[vec![Value::Integer(1), Value::Integer(10)], vec![Value::Integer(2), Value::Null],]
);
}
fn crossed(cross: &CrossProduct, left: &[i32]) -> Vec<Vec<Value>> {
let mut local = cross.local();
let mut chunk = chunk(left);
let mut out = Vec::new();
loop {
let progress = cross.push(&mut chunk, &mut local).expect("a chunk");
out.extend(
(0..chunk.len()).map(|row| vec![chunk.value_at(row, 0), chunk.value_at(row, 1)]),
);
if progress != Progress::Again {
return out;
}
chunk = Chunk::empty(&[]);
}
}
fn kept(memory: &Memory, first: &[i32], second: &[i32]) -> (Keep<'static>, Buffered) {
let (keep, out) = Keep::new(memory);
let mut local = keep.local();
keep.sink(&chunk(first), &mut local).expect("the first right chunk");
keep.sink(&chunk(second), &mut local).expect("the second");
keep.combine(local).expect("the one instance");
keep.finalize(&rudb_pipeline::Lease::alone()).expect("the chunks");
(keep, out)
}
#[test]
fn a_cross_product_pairs_one_left_row_with_one_right_chunk_at_a_time() {
let memory = Memory::unlimited();
let (_keep, right) = kept(&memory, &[10, 20], &[30]);
let cross = CrossProduct::new(&schema("a", 0), &schema("b", 1), right);
assert_eq!(
crossed(&cross, &[1, 2]),
[
vec![Value::Integer(1), Value::Integer(10)],
vec![Value::Integer(1), Value::Integer(20)],
vec![Value::Integer(1), Value::Integer(30)],
vec![Value::Integer(2), Value::Integer(10)],
vec![Value::Integer(2), Value::Integer(20)],
vec![Value::Integer(2), Value::Integer(30)],
]
);
}
#[test]
fn a_cross_product_with_nothing_on_the_right_produces_nothing() {
let memory = Memory::unlimited();
let (keep, right) = Keep::new(&memory);
keep.combine(keep.local()).expect("an instance that saw nothing");
keep.finalize(&rudb_pipeline::Lease::alone()).expect("no chunks");
let cross = CrossProduct::new(&schema("a", 0), &schema("b", 1), right);
let mut local = cross.local();
let mut chunk = chunk(&[1, 2]);
assert_eq!(cross.push(&mut chunk, &mut local).expect("no rows"), Progress::More);
assert!(chunk.is_empty());
}
#[test]
fn an_equality_between_two_columns_of_opposite_sides_is_a_key() {
let mut plan = Plan::new();
let (left, right) = (schema("a", 0), schema("b", 1));
let one = column(&mut plan, 0, LogicalType::Integer);
let other = column(&mut plan, 1, LogicalType::Integer);
let condition = equal(&mut plan, one, other);
let found = equalities(&plan, &[condition], &left, &right).expect("a key");
assert_eq!(found.left, [one]);
assert_eq!(found.right, [other]);
assert_eq!(found.null_is_a_value, [false]);
}
#[test]
fn a_cast_around_one_operand_is_still_a_key() {
let mut plan = Plan::new();
let (left, right) = sides();
let narrow = column(&mut plan, 0, LogicalType::Integer);
let widened =
plan.add_expr(Expr::Cast { input: narrow, try_cast: false }, LogicalType::BigInt);
let other = column(&mut plan, 1, LogicalType::BigInt);
let condition = equal(&mut plan, widened, other);
let found = equalities(&plan, &[condition], &left, &right).expect("a key");
assert_eq!(found.left, [widened]);
assert_eq!(found.right, [other]);
}
#[test]
fn the_gathered_side_written_first_is_lined_back_up() {
let mut plan = Plan::new();
let (left, right) = (schema("a", 0), schema("b", 1));
let driving = column(&mut plan, 0, LogicalType::Integer);
let gathered = column(&mut plan, 1, LogicalType::Integer);
let condition = equal(&mut plan, gathered, driving);
let found = equalities(&plan, &[condition], &left, &right).expect("a key");
assert_eq!(found.left, [driving]);
assert_eq!(found.right, [gathered]);
}
#[test]
fn an_equality_whose_operands_read_one_side_is_not_a_key() {
let mut plan = Plan::new();
let (left, right) = (schema("a", 0), schema("b", 1));
let one = column(&mut plan, 0, LogicalType::Integer);
let condition = equal(&mut plan, one, one);
assert!(equalities(&plan, &[condition], &left, &right).is_none());
}
#[test]
fn an_equality_between_two_constants_is_not_a_key() {
let mut plan = Plan::new();
let (left, right) = (schema("a", 0), schema("b", 1));
let one = plan.add_constant(Value::Integer(1));
let condition = equal(&mut plan, one, one);
assert!(equalities(&plan, &[condition], &left, &right).is_none());
}
#[test]
fn an_operand_that_reads_both_sides_is_not_one_sides_key() {
let mut plan = Plan::new();
let (left, right) = (schema("a", 0), schema("b", 1));
let one = column(&mut plan, 0, LogicalType::Integer);
let other = column(&mut plan, 1, LogicalType::Integer);
let name = plan.intern("+");
let args = plan.add_expr_list(&[one, other]);
let sum = plan.add_expr(Expr::Function { name, args }, LogicalType::Integer);
assert_eq!(side_of(&plan, one, &left, &right), Some(Side::Driving));
assert_eq!(side_of(&plan, other, &left, &right), Some(Side::Gathered));
assert_eq!(side_of(&plan, sum, &left, &right), None);
}
#[test]
fn an_equality_beside_an_inequality_is_still_a_key() {
let mut plan = Plan::new();
let (left, right) = (pair_schema(0), pair_schema(1));
let one = column_at(&mut plan, 0, 0, LogicalType::Integer);
let other = column_at(&mut plan, 1, 0, LogicalType::Integer);
let key = equal(&mut plan, one, other);
let above = column_at(&mut plan, 0, 1, LogicalType::Integer);
let below = column_at(&mut plan, 1, 1, LogicalType::Integer);
let beside = greater(&mut plan, above, below);
let found = equalities(&plan, &[key, beside], &left, &right).expect("a key");
assert_eq!(found.left, [one]);
assert_eq!(found.right, [other]);
assert_eq!(found.residual, [beside]);
}
#[test]
fn a_condition_with_no_equality_in_it_is_not_a_key() {
let mut plan = Plan::new();
let (left, right) = (pair_schema(0), pair_schema(1));
let one = column_at(&mut plan, 0, 0, LogicalType::Integer);
let other = column_at(&mut plan, 1, 0, LogicalType::Integer);
let condition = greater(&mut plan, one, other);
assert!(equalities(&plan, &[condition], &left, &right).is_none());
}
#[test]
fn a_join_with_no_condition_at_all_is_not_a_key() {
let plan = Plan::new();
assert!(equalities(&plan, &[], &schema("a", 0), &schema("b", 1)).is_none());
}
#[test]
fn a_probe_evaluates_the_rest_of_the_condition_on_the_pairs_the_lookup_found() {
let mut plan = Plan::new();
let (left, right) = (pair_schema(0), pair_schema(1));
let key = {
let one = column_at(&mut plan, 0, 0, LogicalType::Integer);
let other = column_at(&mut plan, 1, 0, LogicalType::Integer);
equal(&mut plan, one, other)
};
let beside = {
let one = column_at(&mut plan, 0, 1, LogicalType::Integer);
let other = column_at(&mut plan, 1, 1, LogicalType::Integer);
greater(&mut plan, one, other)
};
let conditions = plan.add_expr_list(&[key, beside]);
let memory = Memory::unlimited();
let (keep, rows) = Keep::new(&memory);
let mut local = keep.local();
keep.sink(&pair_chunk(&[(2, 5), (2, 50), (3, 5)]), &mut local).expect("the gathered rows");
keep.combine(local).expect("the one instance");
keep.finalize(&rudb_pipeline::Lease::alone()).expect("the chunks");
let probe = Probe::new(
&plan,
&left,
&Gathered { schema: &right, chunks: rows, marker: None, swapped: false },
JoinKind::Inner,
conditions,
&Cancel::new(),
&memory,
)
.expect("one equality is enough to look up");
assert_eq!(
probed(&probe, &pair_chunk(&[(1, 10), (2, 10), (3, 10)]), 4),
[
vec![Value::Integer(2), Value::Integer(10), Value::Integer(2), Value::Integer(5)],
vec![Value::Integer(3), Value::Integer(10), Value::Integer(3), Value::Integer(5)],
]
);
}
#[test]
fn a_residual_over_more_pairs_than_fit_in_a_batch_answers_the_same() {
let mut plan = Plan::new();
let (left, right) = (pair_schema(0), pair_schema(1));
let key = {
let one = column_at(&mut plan, 0, 0, LogicalType::Integer);
let other = column_at(&mut plan, 1, 0, LogicalType::Integer);
equal(&mut plan, one, other)
};
let beside = {
let one = column_at(&mut plan, 0, 1, LogicalType::Integer);
let other = column_at(&mut plan, 1, 1, LogicalType::Integer);
greater(&mut plan, one, other)
};
let conditions = plan.add_expr_list(&[key, beside]);
let side: Vec<(i32, i32)> = (0..200).map(|at| (7, at)).collect();
let memory = Memory::unlimited();
let (keep, rows) = Keep::new(&memory);
let mut local = keep.local();
keep.sink(&pair_chunk(&side), &mut local).expect("the gathered rows");
keep.combine(local).expect("the one instance");
keep.finalize(&rudb_pipeline::Lease::alone()).expect("the chunks");
let probe = Probe::new(
&plan,
&left,
&Gathered { schema: &right, chunks: rows, marker: None, swapped: false },
JoinKind::Inner,
conditions,
&Cancel::new(),
&memory,
)
.expect("one equality is enough to look up");
let answer = probed(&probe, &pair_chunk(&side), 4);
assert_eq!(answer.len(), 200 * 199 / 2);
assert!(
answer.iter().all(|row| match (&row[1], &row[3]) {
(Value::Integer(driving), Value::Integer(gathered)) => driving > gathered,
_ => false,
}),
"every pair the residual kept is one it should have"
);
}
#[test]
fn a_left_join_pads_a_driving_row_the_residual_threw_every_candidate_away_for() {
let mut plan = Plan::new();
let (left, right) = (pair_schema(0), pair_schema(1));
let key = {
let one = column_at(&mut plan, 0, 0, LogicalType::Integer);
let other = column_at(&mut plan, 1, 0, LogicalType::Integer);
equal(&mut plan, one, other)
};
let beside = {
let one = column_at(&mut plan, 0, 1, LogicalType::Integer);
let other = column_at(&mut plan, 1, 1, LogicalType::Integer);
greater(&mut plan, one, other)
};
let conditions = plan.add_expr_list(&[key, beside]);
let memory = Memory::unlimited();
let (keep, rows) = Keep::new(&memory);
let mut local = keep.local();
keep.sink(&pair_chunk(&[(2, 50)]), &mut local).expect("the gathered rows");
keep.combine(local).expect("the one instance");
keep.finalize(&rudb_pipeline::Lease::alone()).expect("the chunks");
let probe = Probe::new(
&plan,
&left,
&Gathered { schema: &right, chunks: rows, marker: None, swapped: false },
JoinKind::Left,
conditions,
&Cancel::new(),
&memory,
)
.expect("one equality is enough to look up");
assert_eq!(
probed(&probe, &pair_chunk(&[(2, 10)]), 4),
[vec![Value::Integer(2), Value::Integer(10), Value::Null, Value::Null]]
);
}
fn offered(
plan: &Plan,
left: &Schema,
right: &Schema,
kind: JoinKind,
conditions: Slice,
) -> Option<(ExprRef, ColumnBinding)> {
let memory = Memory::unlimited();
let (keep, rows) = Keep::new(&memory);
let local = keep.local();
keep.combine(local).expect("the one instance");
keep.finalize(&rudb_pipeline::Lease::alone()).expect("the chunks");
let probe = Probe::new(
plan,
left,
&Gathered { schema: right, chunks: rows, marker: None, swapped: false },
kind,
conditions,
&Cancel::new(),
&memory,
)
.expect("one equality is enough to look up");
probe.sideways()
}
#[test]
fn an_inner_join_on_a_column_offers_its_key_to_the_scan_below() {
let mut plan = Plan::new();
let (left, right) = (pair_schema(0), pair_schema(1));
let one = column_at(&mut plan, 0, 0, LogicalType::Integer);
let other = column_at(&mut plan, 1, 0, LogicalType::Integer);
let condition = equal(&mut plan, one, other);
let conditions = plan.add_expr_list(&[condition]);
assert_eq!(
offered(&plan, &left, &right, JoinKind::Inner, conditions),
Some((other, ColumnBinding::new(0, 0)))
);
assert_eq!(
offered(&plan, &left, &right, JoinKind::Semi, conditions),
Some((other, ColumnBinding::new(0, 0))),
"a semi join drops an unmatched driving row too"
);
}
#[test]
fn a_join_that_keeps_an_unmatched_driving_row_offers_nothing() {
let mut plan = Plan::new();
let (left, right) = (pair_schema(0), pair_schema(1));
let one = column_at(&mut plan, 0, 0, LogicalType::Integer);
let other = column_at(&mut plan, 1, 0, LogicalType::Integer);
let condition = equal(&mut plan, one, other);
let conditions = plan.add_expr_list(&[condition]);
for kind in [JoinKind::Left, JoinKind::Anti, JoinKind::Single] {
assert_eq!(offered(&plan, &left, &right, kind, conditions), None, "{kind:?}");
}
}
#[test]
fn an_equality_that_matches_two_nulls_offers_nothing() {
let mut plan = Plan::new();
let (left, right) = (pair_schema(0), pair_schema(1));
let one = column_at(&mut plan, 0, 0, LogicalType::Integer);
let other = column_at(&mut plan, 1, 0, LogicalType::Integer);
let condition = plan.add_expr(
Expr::Compare { op: CompareOp::NotDistinctFrom, left: one, right: other },
LogicalType::Boolean,
);
let conditions = plan.add_expr_list(&[condition]);
assert_eq!(offered(&plan, &left, &right, JoinKind::Inner, conditions), None);
}
#[test]
fn a_driving_key_that_is_an_expression_offers_nothing() {
let mut plan = Plan::new();
let (left, right) = sides();
let narrow = column(&mut plan, 0, LogicalType::Integer);
let widened =
plan.add_expr(Expr::Cast { input: narrow, try_cast: false }, LogicalType::BigInt);
let other = column(&mut plan, 1, LogicalType::BigInt);
let condition = equal(&mut plan, widened, other);
let conditions = plan.add_expr_list(&[condition]);
assert_eq!(offered(&plan, &left, &right, JoinKind::Inner, conditions), None);
}
#[test]
fn a_full_join_reports_a_gathered_row_whose_only_candidate_the_residual_threw_away() {
let mut plan = Plan::new();
let (left, right) = (pair_schema(0), pair_schema(1));
let key = {
let one = column_at(&mut plan, 0, 0, LogicalType::Integer);
let other = column_at(&mut plan, 1, 0, LogicalType::Integer);
equal(&mut plan, one, other)
};
let beside = {
let one = column_at(&mut plan, 0, 1, LogicalType::Integer);
let other = column_at(&mut plan, 1, 1, LogicalType::Integer);
greater(&mut plan, one, other)
};
let conditions = plan.add_expr_list(&[key, beside]);
let memory = Memory::unlimited();
let (keep, gathered) = Keep::new(&memory);
let mut local = keep.local();
keep.sink(&pair_chunk(&[(2, 50), (4, 1)]), &mut local).expect("the gathered rows");
keep.combine(local).expect("the one instance");
keep.finalize(&rudb_pipeline::Lease::alone()).expect("the chunks");
let (join, out) = Join::new(
&plan,
&left,
Gathered { schema: &right, chunks: gathered, marker: None, swapped: false },
JoinKind::Full,
conditions,
&Cancel::new(),
&memory,
);
let mut local = join.local();
join.sink(&pair_chunk(&[(2, 10)]), &mut local).expect("the driving rows");
join.combine(local).expect("the one instance");
join.finalize(&rudb_pipeline::Lease::alone()).expect("the answer");
assert_eq!(
rows(&out, 4),
[
vec![Value::Integer(2), Value::Integer(10), Value::Null, Value::Null],
vec![Value::Null, Value::Null, Value::Integer(2), Value::Integer(50)],
vec![Value::Null, Value::Null, Value::Integer(4), Value::Integer(1)],
]
);
}
#[test]
fn a_probe_answers_a_join_whose_key_is_a_cast() {
let mut plan = Plan::new();
let (left, right) = sides();
let narrow = column(&mut plan, 0, LogicalType::Integer);
let widened =
plan.add_expr(Expr::Cast { input: narrow, try_cast: false }, LogicalType::BigInt);
let other = column(&mut plan, 1, LogicalType::BigInt);
let condition = equal(&mut plan, widened, other);
let conditions = plan.add_expr_list(&[condition]);
let memory = Memory::unlimited();
let (keep, rows) = Keep::new(&memory);
let mut local = keep.local();
keep.sink(&wide_chunk(&[2, 3, 4]), &mut local).expect("the gathered rows");
keep.combine(local).expect("the one instance");
keep.finalize(&rudb_pipeline::Lease::alone()).expect("the chunks");
let probe = Probe::new(
&plan,
&left,
&Gathered { schema: &right, chunks: rows, marker: None, swapped: false },
JoinKind::Inner,
conditions,
&Cancel::new(),
&memory,
)
.expect("a lookup answers an inner join on one equality");
assert_eq!(
probed(&probe, &chunk(&[1, 2, 3]), 2),
[vec![Value::Integer(2), Value::BigInt(2)], vec![Value::Integer(3), Value::BigInt(3)]]
);
}
fn subject(rows: &[Option<i32>], memory: &Memory) -> Buffered {
let (keep, kept) = Keep::new(memory);
let mut local = keep.local();
if !rows.is_empty() {
keep.sink(&some_chunk(rows), &mut local).expect("the gathered rows");
}
keep.combine(local).expect("the one instance");
keep.finalize(&rudb_pipeline::Lease::alone()).expect("the chunks");
kept
}
fn marking(
kind: JoinKind,
subject_rows: &[Option<i32>],
driving: &[&[Option<i32>]],
) -> Vec<Value> {
let mut plan = Plan::new();
let (left, right) = (schema("b", 1), schema("a", 0));
let conditions = {
let one = column(&mut plan, 1, LogicalType::Integer);
let other = column(&mut plan, 0, LogicalType::Integer);
let key = equal(&mut plan, one, other);
plan.add_expr_list(&[key])
};
let memory = Memory::unlimited();
let kept = subject(subject_rows, &memory);
let (mark, out) = Marking::new(
&plan,
&left,
&Gathered { schema: &right, chunks: kept, marker: None, swapped: true },
kind,
conditions,
&Cancel::new(),
&memory,
)
.expect("one equality is enough to mark on");
for rows in driving {
let mut local = mark.local();
mark.sink(&some_chunk(rows), &mut local).expect("a driving chunk");
mark.combine(local).expect("one instance");
}
mark.finalize(&rudb_pipeline::Lease::alone()).expect("the answer");
answered_rows(&out)
}
fn answered_rows(out: &Buffered) -> Vec<Value> {
let reader = out.reader();
let mut rows = Vec::new();
for at in 0..reader.len().expect("the chunks") {
let chunk = reader.at(at).expect("the chunk").expect("a chunk that was counted");
rows.extend((0..chunk.len()).map(|row| chunk.value_at(row, 0)));
}
rows
}
#[test]
fn a_marking_semi_join_answers_with_the_gathered_rows_something_matched() {
assert_eq!(
marking(JoinKind::Semi, &[Some(1), Some(2), Some(3)], &[&[Some(3), Some(1)]]),
[Value::Integer(1), Value::Integer(3)]
);
}
#[test]
fn a_marking_anti_join_answers_with_the_gathered_rows_nothing_matched() {
assert_eq!(
marking(JoinKind::Anti, &[Some(1), Some(2), Some(3)], &[&[Some(3), Some(1)]]),
[Value::Integer(2)]
);
}
#[test]
fn a_marking_semi_join_answers_a_subject_row_once_however_often_it_matched() {
assert_eq!(
marking(JoinKind::Semi, &[Some(1), Some(2)], &[&[Some(2), Some(2), Some(2)]]),
[Value::Integer(2)]
);
}
#[test]
fn a_marking_join_puts_the_bits_of_two_instances_together() {
assert_eq!(
marking(JoinKind::Semi, &[Some(1), Some(2), Some(3)], &[&[Some(1)], &[Some(3)]]),
[Value::Integer(1), Value::Integer(3)]
);
}
#[test]
fn a_marking_join_over_a_driving_side_with_no_rows_marks_nothing() {
assert_eq!(marking(JoinKind::Semi, &[Some(1), Some(2)], &[]), []);
assert_eq!(
marking(JoinKind::Anti, &[Some(1), Some(2)], &[]),
[Value::Integer(1), Value::Integer(2)]
);
}
#[test]
fn a_marking_join_never_marks_a_subject_row_whose_key_is_null() {
assert_eq!(
marking(JoinKind::Semi, &[Some(1), None], &[&[Some(1), None]]),
[Value::Integer(1)]
);
assert_eq!(marking(JoinKind::Anti, &[Some(1), None], &[&[Some(1), None]]), [Value::Null]);
}
#[test]
fn a_marking_join_over_an_empty_subject_answers_nothing() {
assert_eq!(marking(JoinKind::Semi, &[], &[&[Some(1)]]), []);
assert_eq!(marking(JoinKind::Anti, &[], &[&[Some(1)]]), []);
}
#[test]
fn a_marking_join_marks_only_what_the_residual_kept() {
let mut plan = Plan::new();
let (left, right) = (pair_schema(1), pair_schema(0));
let conditions = {
let one = column_at(&mut plan, 1, 0, LogicalType::Integer);
let other = column_at(&mut plan, 0, 0, LogicalType::Integer);
let key = equal(&mut plan, one, other);
let driving_g = column_at(&mut plan, 1, 1, LogicalType::Integer);
let subject_g = column_at(&mut plan, 0, 1, LogicalType::Integer);
let over = greater(&mut plan, driving_g, subject_g);
plan.add_expr_list(&[key, over])
};
let memory = Memory::unlimited();
let (keep, kept) = Keep::new(&memory);
let mut local = keep.local();
keep.sink(&pair_chunk(&[(7, 1), (7, 9), (8, 1)]), &mut local).expect("the gathered rows");
keep.combine(local).expect("the one instance");
keep.finalize(&rudb_pipeline::Lease::alone()).expect("the chunks");
let (mark, out) = Marking::new(
&plan,
&left,
&Gathered { schema: &right, chunks: kept, marker: None, swapped: true },
JoinKind::Semi,
conditions,
&Cancel::new(),
&memory,
)
.expect("an equality beside a residual is still a lookup");
let mut instance = mark.local();
mark.sink(&pair_chunk(&[(7, 5), (8, 0)]), &mut instance).expect("a driving chunk");
mark.combine(instance).expect("the one instance");
mark.finalize(&rudb_pipeline::Lease::alone()).expect("the answer");
assert_eq!(answered_rows(&out), [Value::Integer(7)]);
}
#[test]
fn a_marking_join_offers_its_key_to_the_scan_under_either_kind() {
for kind in [JoinKind::Semi, JoinKind::Anti] {
let mut plan = Plan::new();
let (left, right) = (schema("b", 1), schema("a", 0));
let conditions = {
let one = column(&mut plan, 1, LogicalType::Integer);
let other = column(&mut plan, 0, LogicalType::Integer);
let key = equal(&mut plan, one, other);
plan.add_expr_list(&[key])
};
let memory = Memory::unlimited();
let kept = subject(&[Some(1)], &memory);
let (mark, _out) = Marking::new(
&plan,
&left,
&Gathered { schema: &right, chunks: kept, marker: None, swapped: true },
kind,
conditions,
&Cancel::new(),
&memory,
)
.expect("one equality is enough to mark on");
assert!(mark.sideways().is_some(), "a {} join has a key to offer", kind.keyword());
}
}
#[test]
fn a_marking_join_refuses_a_kind_it_does_not_answer() {
let mut plan = Plan::new();
let (left, right) = (schema("b", 1), schema("a", 0));
let conditions = {
let one = column(&mut plan, 1, LogicalType::Integer);
let other = column(&mut plan, 0, LogicalType::Integer);
let key = equal(&mut plan, one, other);
plan.add_expr_list(&[key])
};
let memory = Memory::unlimited();
for kind in [JoinKind::Inner, JoinKind::Left, JoinKind::Single, JoinKind::Mark] {
let kept = subject(&[Some(1)], &memory);
let made = Marking::new(
&plan,
&left,
&Gathered { schema: &right, chunks: kept, marker: None, swapped: true },
kind,
conditions,
&Cancel::new(),
&memory,
);
assert!(made.is_none(), "a {} join is not a marking join", kind.keyword());
}
}
#[derive(Debug)]
struct Answered {
pairs: Vec<Vec<Value>>,
padded: Vec<Vec<Value>>,
chunks: Vec<usize>,
}
fn rows_of(chunk: &Chunk, width: usize) -> Vec<Vec<Value>> {
(0..chunk.len())
.map(|row| (0..width).map(|column| chunk.value_at(row, column)).collect())
.collect()
}
fn padding(kind: JoinKind, gathered: &[Option<i32>], driving: &[&[Option<i32>]]) -> Answered {
let mut plan = Plan::new();
let (left, right) = (schema("b", 1), schema("a", 0));
let conditions = {
let one = column(&mut plan, 1, LogicalType::Integer);
let other = column(&mut plan, 0, LogicalType::Integer);
let key = equal(&mut plan, one, other);
plan.add_expr_list(&[key])
};
let memory = Memory::unlimited();
let kept = gathered_side(gathered, &memory);
let pad = Padding::new(
&plan,
&left,
&Gathered { schema: &right, chunks: kept, marker: None, swapped: true },
kind,
conditions,
&Cancel::new(),
&memory,
)
.expect("one equality is enough to pair on");
answered(&pad, &driving.iter().map(|rows| some_chunk(rows)).collect::<Vec<Chunk>>(), 2)
}
fn gathered_side(rows: &[Option<i32>], memory: &Memory) -> Buffered {
let (keep, kept) = Keep::new(memory);
let mut local = keep.local();
for piece in rows.chunks(VECTOR_SIZE) {
keep.sink(&some_chunk(piece), &mut local).expect("the gathered rows");
}
keep.combine(local).expect("the one instance");
keep.finalize(&rudb_pipeline::Lease::alone()).expect("the chunks");
kept
}
fn answered(pad: &Padding<'_>, driving: &[Chunk], width: usize) -> Answered {
pad.prepare(&rudb_pipeline::Lease::alone()).expect("the table");
let mut pairs = Vec::new();
for chunk in driving {
let mut local = pad.local();
let mut chunk = chunk.clone();
loop {
let progress = pad.push(&mut chunk, &mut local).expect("a chunk");
pairs.extend(rows_of(&chunk, width));
if progress != Progress::Again {
break;
}
chunk = Chunk::empty(&[]);
}
}
let mut padded = Vec::new();
let mut chunks = Vec::new();
pad.drain(&mut |chunk| {
chunks.push(chunk.len());
padded.extend(rows_of(chunk, width));
Ok(Progress::More)
})
.expect("the gathered rows nothing matched");
Answered { pairs, padded, chunks }
}
#[test]
fn a_padding_right_join_streams_the_pairs_and_pads_what_nothing_matched() {
let answer = padding(JoinKind::Right, &[Some(1), Some(2), Some(3)], &[&[Some(3), Some(1)]]);
assert_eq!(
answer.pairs,
[
vec![Value::Integer(3), Value::Integer(3)],
vec![Value::Integer(1), Value::Integer(1)]
]
);
assert_eq!(answer.padded, [vec![Value::Integer(2), Value::Null]]);
}
#[test]
fn a_padding_full_join_pads_the_driving_side_as_it_goes_and_the_gathered_side_at_the_end() {
let answer = padding(JoinKind::Full, &[Some(1), Some(2)], &[&[Some(2), Some(9)]]);
assert_eq!(
answer.pairs,
[vec![Value::Integer(2), Value::Integer(2)], vec![Value::Null, Value::Integer(9)]]
);
assert_eq!(answer.padded, [vec![Value::Integer(1), Value::Null]]);
}
#[test]
fn a_padding_join_pads_a_gathered_row_no_times_however_often_it_matched() {
let answer = padding(JoinKind::Right, &[Some(1), Some(2)], &[&[Some(2), Some(2), Some(2)]]);
assert_eq!(answer.pairs.len(), 3);
assert_eq!(answer.padded, [vec![Value::Integer(1), Value::Null]]);
}
#[test]
fn a_padding_join_puts_the_bits_of_two_instances_together() {
let answer =
padding(JoinKind::Right, &[Some(1), Some(2), Some(3)], &[&[Some(1)], &[Some(3)]]);
assert_eq!(answer.padded, [vec![Value::Integer(2), Value::Null]]);
}
#[test]
fn a_padding_join_over_a_driving_side_with_no_rows_pads_every_gathered_row() {
let answer = padding(JoinKind::Right, &[Some(1), Some(2)], &[]);
assert_eq!(answer.pairs, Vec::<Vec<Value>>::new());
assert_eq!(
answer.padded,
[vec![Value::Integer(1), Value::Null], vec![Value::Integer(2), Value::Null]]
);
}
#[test]
fn a_padding_join_over_an_empty_gathered_side_owes_nothing() {
let right = padding(JoinKind::Right, &[], &[&[Some(1)]]);
assert_eq!(right.pairs, Vec::<Vec<Value>>::new());
assert_eq!(right.padded, Vec::<Vec<Value>>::new());
let full = padding(JoinKind::Full, &[], &[&[Some(1)]]);
assert_eq!(full.pairs, [vec![Value::Null, Value::Integer(1)]]);
assert_eq!(full.padded, Vec::<Vec<Value>>::new());
}
#[test]
fn a_padding_join_never_matches_a_key_that_is_null() {
let answer = padding(JoinKind::Right, &[Some(1), None], &[&[Some(1), None]]);
assert_eq!(answer.pairs, [vec![Value::Integer(1), Value::Integer(1)]]);
assert_eq!(answer.padded, [vec![Value::Null, Value::Null]]);
}
#[test]
fn a_padding_join_pads_a_gathered_row_the_residual_threw_away() {
let mut plan = Plan::new();
let (left, right) = (pair_schema(1), pair_schema(0));
let conditions = {
let one = column_at(&mut plan, 1, 0, LogicalType::Integer);
let other = column_at(&mut plan, 0, 0, LogicalType::Integer);
let key = equal(&mut plan, one, other);
let driving_g = column_at(&mut plan, 1, 1, LogicalType::Integer);
let gathered_g = column_at(&mut plan, 0, 1, LogicalType::Integer);
let over = greater(&mut plan, driving_g, gathered_g);
plan.add_expr_list(&[key, over])
};
let memory = Memory::unlimited();
let (keep, kept) = Keep::new(&memory);
let mut local = keep.local();
keep.sink(&pair_chunk(&[(7, 1), (7, 9), (8, 1)]), &mut local).expect("the gathered rows");
keep.combine(local).expect("the one instance");
keep.finalize(&rudb_pipeline::Lease::alone()).expect("the chunks");
let pad = Padding::new(
&plan,
&left,
&Gathered { schema: &right, chunks: kept, marker: None, swapped: true },
JoinKind::Right,
conditions,
&Cancel::new(),
&memory,
)
.expect("an equality beside a residual is still a lookup");
let answer = answered(&pad, &[pair_chunk(&[(7, 5), (8, 0)])], 4);
assert_eq!(
answer.pairs,
[vec![Value::Integer(7), Value::Integer(1), Value::Integer(7), Value::Integer(5)]]
);
assert_eq!(
answer.padded,
[
vec![Value::Integer(7), Value::Integer(9), Value::Null, Value::Null],
vec![Value::Integer(8), Value::Integer(1), Value::Null, Value::Null]
]
);
}
#[test]
fn a_padding_joins_drain_hands_over_a_vector_at_a_time() {
let rows: Vec<Option<i32>> = (0..VECTOR_SIZE as i32 + 5).map(Some).collect();
let answer = padding(JoinKind::Right, &rows, &[]);
assert_eq!(answer.chunks, [VECTOR_SIZE, 5]);
assert_eq!(answer.padded.len(), VECTOR_SIZE + 5);
}
#[test]
fn a_padding_join_offers_no_key_to_the_scan() {
let mut plan = Plan::new();
let (left, right) = (schema("b", 1), schema("a", 0));
let conditions = {
let one = column(&mut plan, 1, LogicalType::Integer);
let other = column(&mut plan, 0, LogicalType::Integer);
let key = equal(&mut plan, one, other);
plan.add_expr_list(&[key])
};
let memory = Memory::unlimited();
let kept = subject(&[Some(1)], &memory);
let pad = Padding::new(
&plan,
&left,
&Gathered { schema: &right, chunks: kept, marker: None, swapped: true },
JoinKind::Right,
conditions,
&Cancel::new(),
&memory,
)
.expect("one equality is enough to pair on");
assert!(pad.sideways().is_none(), "a padding join has no key it can offer");
}
#[test]
fn a_padding_join_refuses_a_kind_it_does_not_answer() {
let mut plan = Plan::new();
let (left, right) = (schema("b", 1), schema("a", 0));
let conditions = {
let one = column(&mut plan, 1, LogicalType::Integer);
let other = column(&mut plan, 0, LogicalType::Integer);
let key = equal(&mut plan, one, other);
plan.add_expr_list(&[key])
};
let memory = Memory::unlimited();
for kind in
[JoinKind::Inner, JoinKind::Left, JoinKind::Semi, JoinKind::Anti, JoinKind::Mark]
{
let kept = subject(&[Some(1)], &memory);
let made = Padding::new(
&plan,
&left,
&Gathered { schema: &right, chunks: kept, marker: None, swapped: true },
kind,
conditions,
&Cancel::new(),
&memory,
);
assert!(made.is_none(), "a {} join is not a padding join", kind.keyword());
}
}
}