use std::sync::atomic::{AtomicU32, Ordering};
use rudb_common::{Cancel, Error, LogicalType, Result};
use rudb_pipeline::Lease;
use rudb_vector::{Form, Vector};
use crate::pairs::in_parallel;
use crate::table::{BATCH, Probe, Table, Walk};
const NONE: u32 = u32::MAX;
const SPLIT: usize = 64 * 1024;
pub(crate) const MISS: usize = usize::MAX;
#[derive(Debug)]
struct Part {
table: Table,
base: usize,
}
#[derive(Debug, Default)]
pub(crate) struct Lookup {
parts: Vec<Part>,
bits: u32,
head: Vec<u32>,
next: Vec<AtomicU32>,
kept: usize,
}
impl Lookup {
pub(crate) fn build(
keys: &[Vector],
rows: usize,
nulls: &[bool],
threads: &Lease<'_>,
cancel: &Cancel,
) -> Result<Self> {
if rows >= NONE as usize {
return Err(too_many_rows());
}
if rows == 0 || keys.is_empty() {
return Ok(Self::default());
}
let mut hashes = Vec::new();
crate::table::hash(keys, rows, &mut hashes, crate::table::Across::TwoInputs);
let mut keyed = Vec::new();
which_are_keyed(keys, rows, nulls, &mut keyed);
let bits = split_into(rows, threads.degree());
let count = 1usize << bits;
let next: Vec<AtomicU32> = (0..rows).map(|_| AtomicU32::new(NONE)).collect();
let types: Vec<LogicalType> = keys.iter().map(|key| key.logical_type().clone()).collect();
let one = |part: usize| -> Result<(Table, Vec<u32>, usize)> {
fill(part, bits, &types, keys, &hashes, &keyed, &next, cancel)
};
let filled = in_parallel(threads, count, threads.degree(), "join table partition", one)?;
let mut parts = Vec::with_capacity(count);
let mut head = Vec::new();
let mut kept = 0;
for (table, mine, held) in filled {
parts.push(Part { base: head.len(), table });
head.extend(mine);
kept += held;
}
Ok(Self { parts, bits, head, next, kept })
}
pub(crate) fn is_empty(&self) -> bool {
self.kept == 0
}
pub(crate) fn footprint(&self) -> u64 {
let tables: u64 =
self.parts.iter().map(|part| part.table.footprint() + part.table.owned()).sum();
let chain =
self.head.capacity() * size_of::<u32>() + self.next.capacity() * size_of::<AtomicU32>();
tables + u64::try_from(chain).unwrap_or(u64::MAX)
}
pub(crate) fn slots(
&self,
keys: &[Vector],
rows: usize,
nulls: &[bool],
scratch: &mut Scratch,
into: &mut Vec<usize>,
) {
into.clear();
into.resize(rows, MISS);
if rows == 0 || self.parts.is_empty() {
return;
}
crate::table::hash(keys, rows, &mut scratch.hashes, crate::table::Across::TwoInputs);
which_are_keyed(keys, rows, nulls, &mut scratch.keyed);
if self.parts.len() == 1 {
let mut from = 0;
while from < rows {
let upto = (from + BATCH).min(rows);
self.parts[0].table.probe_run(
&scratch.hashes,
keys,
from,
upto,
into,
&mut scratch.walk,
);
from = upto;
}
} else {
self.deal(keys, rows, scratch, into);
}
for (row, &keyed) in scratch.keyed.iter().enumerate().take(rows) {
if !keyed {
into[row] = MISS;
}
}
}
fn deal(&self, keys: &[Vector], rows: usize, scratch: &mut Scratch, into: &mut [usize]) {
scratch.by_part.resize_with(self.parts.len(), Vec::new);
for held in &mut scratch.by_part {
held.clear();
}
for row in 0..rows {
if scratch.keyed[row] {
scratch.by_part[part_of(scratch.hashes[row], self.bits)].push(row);
}
}
for (part, mine) in self.parts.iter().zip(&scratch.by_part) {
let mut from = 0;
while from < mine.len() {
let upto = (from + BATCH).min(mine.len());
let batch = &mine[from..upto];
scratch.found.clear();
scratch.found.resize(batch.len(), MISS);
part.table.probe_these(
&scratch.hashes,
keys,
batch,
&mut scratch.found,
&mut scratch.walk,
);
for (&row, &slot) in batch.iter().zip(&scratch.found) {
if slot != MISS {
into[row] = part.base + slot;
}
}
from = upto;
}
}
}
pub(crate) fn matches(&self, slot: usize, into: &mut Vec<u32>) {
into.clear();
if slot == MISS {
return;
}
let mut at = self.head[slot];
while at != NONE {
into.push(at);
at = self.next[at as usize].load(Ordering::Relaxed);
}
}
}
#[allow(clippy::too_many_arguments)]
fn fill(
part: usize,
bits: u32,
types: &[LogicalType],
keys: &[Vector],
hashes: &[u64],
keyed: &[bool],
next: &[AtomicU32],
cancel: &Cancel,
) -> Result<(Table, Vec<u32>, usize)> {
let mut mine: Vec<usize> = Vec::new();
for (row, &hash) in hashes.iter().enumerate() {
if keyed[row] && part_of(hash, bits) == part {
mine.push(row);
}
}
let mut table = Table::new(types);
let mut head: Vec<u32> = Vec::new();
let mut tail: Vec<u32> = Vec::new();
let mut found: Vec<usize> = Vec::new();
let mut walk = Walk::default();
let mut kept = 0;
let mut from = 0;
while from < mine.len() {
cancel.check()?;
let upto = (from + BATCH).min(mine.len());
let batch = &mine[from..upto];
found.clear();
found.resize(batch.len(), MISS);
table.probe_these(hashes, keys, batch, &mut found, &mut walk);
for &place in walk.pending() {
let row = batch[place];
match table.probe(hashes[row], keys, row) {
Probe::Found(slot) => found[place] = slot,
Probe::Vacant(bucket) => {
let slot = table.insert(bucket, hashes[row], keys, row)?;
debug_assert_eq!(slot, head.len(), "a slot is the number of keys before it");
head.push(NONE);
tail.push(NONE);
found[place] = slot;
}
}
}
for (&row, &slot) in batch.iter().zip(&found) {
let at = u32::try_from(row).map_err(|_| too_many_rows())?;
if tail[slot] == NONE {
head[slot] = at;
} else {
next[tail[slot] as usize].store(at, Ordering::Relaxed);
}
tail[slot] = at;
kept += 1;
}
from = upto;
}
Ok((table, head, kept))
}
fn split_into(rows: usize, threads: usize) -> u32 {
if rows < SPLIT || threads <= 1 {
return 0;
}
threads.next_power_of_two().trailing_zeros()
}
fn part_of(hash: u64, bits: u32) -> usize {
if bits == 0 {
return 0;
}
(hash >> (64 - bits)) as usize
}
#[derive(Debug, Default)]
pub(crate) struct Scratch {
hashes: Vec<u64>,
keyed: Vec<bool>,
walk: Walk,
by_part: Vec<Vec<usize>>,
found: Vec<usize>,
}
fn which_are_keyed(keys: &[Vector], rows: usize, nulls: &[bool], keyed: &mut Vec<bool>) {
keyed.clear();
keyed.resize(rows, true);
for (column, &stored) in keys.iter().zip(nulls) {
if stored || !has_nulls(column, rows) {
continue;
}
for (row, flag) in keyed.iter_mut().enumerate().take(rows) {
*flag = *flag && !column.is_null_at(row);
}
}
}
pub(crate) fn has_nulls(column: &Vector, rows: usize) -> bool {
match column.form() {
Form::Dictionary | Form::Rle => true,
_ => column.validity().has_nulls(rows),
}
}
fn too_many_rows() -> Error {
Error::out_of_memory(format!(
"a hash join cannot gather more than {} rows on one side",
NONE - 1
))
}
#[cfg(test)]
fn column(values: &[Option<i32>]) -> Vector {
use rudb_common::Value;
let values: Vec<Value> =
values.iter().map(|value| value.map_or(Value::Null, Value::Integer)).collect();
Vector::from_values(LogicalType::Integer, &values).expect("a column of integers")
}
#[cfg(test)]
mod tests {
use rudb_common::Cancel;
use rudb_pipeline::{Lease, Pool};
use super::{Lookup, MISS, SPLIT, Scratch, column, part_of, split_into};
fn built(values: &[Option<i32>]) -> Lookup {
built_by(values, &[false], &Lease::alone())
}
fn built_by(values: &[Option<i32>], nulls: &[bool], threads: &Lease<'_>) -> Lookup {
Lookup::build(&[column(values)], values.len(), nulls, threads, &Cancel::new())
.expect("a build")
}
fn found(lookup: &Lookup, values: &[Option<i32>]) -> Vec<Vec<u32>> {
let mut scratch = Scratch::default();
let mut slots = Vec::new();
lookup.slots(&[column(values)], values.len(), &[false], &mut scratch, &mut slots);
let mut chain = Vec::new();
slots
.iter()
.map(|&slot| {
lookup.matches(slot, &mut chain);
chain.clone()
})
.collect()
}
#[test]
fn a_key_with_no_rows_is_a_miss_and_a_key_with_one_is_that_row() {
let lookup = built(&[Some(10), Some(20)]);
assert_eq!(
found(&lookup, &[Some(20), Some(30), Some(10)]),
vec![vec![1], Vec::new(), vec![0]]
);
}
#[test]
fn a_keys_rows_come_out_in_the_order_the_gathered_side_holds_them() {
let lookup = built(&[Some(7), Some(9), Some(7), Some(7), Some(9)]);
assert_eq!(found(&lookup, &[Some(7), Some(9)]), vec![vec![0, 2, 3], vec![1, 4]]);
}
#[test]
fn a_key_first_seen_twice_inside_one_batch_is_one_key() {
let lookup = built(&[Some(4), Some(4)]);
assert_eq!(found(&lookup, &[Some(4)]), vec![vec![0, 1]]);
}
#[test]
fn a_chain_that_spans_several_batches_stays_in_order() {
let values: Vec<Option<i32>> = (0..500).map(|row| Some(row % 3)).collect();
let lookup = built(&values);
let mut chain = Vec::new();
let mut scratch = Scratch::default();
let mut slots = Vec::new();
lookup.slots(&[column(&[Some(1)])], 1, &[false], &mut scratch, &mut slots);
lookup.matches(slots[0], &mut chain);
let wanted: Vec<u32> = (0..500).filter(|row| row % 3 == 1).collect();
assert_eq!(chain, wanted);
}
#[test]
fn a_rejected_null_is_neither_stored_nor_found() {
let lookup = built(&[Some(1), None, Some(2)]);
let mut scratch = Scratch::default();
let mut slots = Vec::new();
let driving = [Some(1), None];
lookup.slots(&[column(&driving)], 2, &[false], &mut scratch, &mut slots);
assert_ne!(slots[0], MISS, "a driving row with a key finds it");
assert_eq!(slots[1], MISS, "a driving row whose key is null finds nothing");
}
#[test]
fn a_null_a_join_calls_a_value_is_stored_and_found() {
let lookup = built_by(&[Some(1), None, None], &[true], &Lease::alone());
let mut scratch = Scratch::default();
let mut slots = Vec::new();
let driving = [None];
lookup.slots(&[column(&driving)], 1, &[true], &mut scratch, &mut slots);
let mut chain = Vec::new();
lookup.matches(slots[0], &mut chain);
assert_eq!(chain, vec![1, 2]);
}
#[test]
fn a_gathered_side_with_nothing_keyed_in_it_is_empty() {
let nothing = Lookup::build(&[], 0, &[false], &Lease::alone(), &Cancel::new())
.expect("no columns at all");
assert!(nothing.is_empty());
assert!(built(&[None, None]).is_empty(), "every row's key was a rejected null");
assert!(!built(&[Some(1)]).is_empty());
}
#[test]
fn a_partition_is_named_by_the_top_bits_and_a_bucket_by_the_low_ones() {
assert_eq!(part_of(0, 2), 0);
assert_eq!(part_of(u64::MAX, 2), 3);
assert_eq!(part_of(1 << 62, 2), 1);
assert_eq!(part_of(u64::MAX, 0), 0, "one partition holds everything");
assert_eq!(part_of(0xFFFF_FFFF, 2), 0, "the low bits say nothing about which partition");
}
#[test]
fn a_small_side_is_not_split_at_all() {
assert_eq!(split_into(1_000, 8), 0);
assert_eq!(split_into(SPLIT - 1, 8), 0);
assert_eq!(split_into(SPLIT, 1), 0);
assert_eq!(split_into(SPLIT, 8), 3);
assert_eq!(split_into(SPLIT, 6), 3, "rounded up to a power of two");
}
#[test]
fn a_side_built_in_partitions_answers_the_same_as_one_built_whole() {
let values: Vec<Option<i32>> = (0..SPLIT as i32 + 1_000).map(|row| Some(row % 7)).collect();
let pool = Pool::new(4);
let lookup = built_by(&values, &[false], &pool.lease(4));
assert!(lookup.parts.len() > 1, "a side this long is split");
let mut scratch = Scratch::default();
let mut slots = Vec::new();
let driving: Vec<Option<i32>> = (0..9).map(Some).collect();
lookup.slots(&[column(&driving)], driving.len(), &[false], &mut scratch, &mut slots);
let mut chain = Vec::new();
for (key, &slot) in slots.iter().enumerate() {
let key = i32::try_from(key).expect("nine of them");
let wanted: Vec<u32> = (0..values.len())
.filter(|&row| values[row] == Some(key))
.map(|row| u32::try_from(row).expect("a side this long"))
.collect();
if wanted.is_empty() {
assert_eq!(slot, MISS, "key {key} is not in the side");
continue;
}
lookup.matches(slot, &mut chain);
assert_eq!(chain, wanted, "key {key} came out in the wrong order");
}
}
}